From 72cbe9e914c34e75ac99d2c67acfe15ebf377bd1 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 01:49:05 -0400 Subject: [PATCH 01/27] cuda : record HC affine fusion source prerequisites Carry the reviewed kernel sources from PR125, PR129 and PR132 as the matching control for the separate tensor-affine HC follow-up. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 655 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/mmvf.cu | 185 ++++++++- ggml/src/ggml-cuda/mmvf.cuh | 4 + ggml/src/ggml-cuda/norm.cu | 340 +++++++++++++++++ ggml/src/ggml-cuda/norm.cuh | 6 + ggml/src/ggml-cuda/unary.cu | 102 +++++ ggml/src/ggml-cuda/unary.cuh | 2 + 7 files changed, 1286 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index f9fb46c2a3b..a7687b61c92 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3445,6 +3445,580 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } +struct ggml_cuda_hc_post_norm_match { + ggml_tensor * post = nullptr; + ggml_tensor * norm = nullptr; + ggml_tensor * mul = nullptr; + int count = 0; +}; + +static ggml_cuda_hc_post_norm_match ggml_cuda_match_hc_post_norm(ggml_cgraph * cgraph, int i) { + ggml_cuda_hc_post_norm_match match; + if (cgraph->nodes[i]->op != GGML_OP_DSV4_HC_POST || i + 1 >= cgraph->n_nodes) { + return match; + } + ggml_tensor * post = cgraph->nodes[i]; + int norm_i = i + 1; + const bool reshape = cgraph->nodes[norm_i]->op == GGML_OP_RESHAPE; + if (reshape) { + const ggml_tensor * view = cgraph->nodes[norm_i]; + if (view->src[0] != post || view->view_src != post || view->view_offs != 0 || !ggml_is_contiguous(view) || ++norm_i >= cgraph->n_nodes) { + return match; + } + } + ggml_tensor * norm = cgraph->nodes[norm_i]; + if (norm->op != GGML_OP_RMS_NORM || norm->src[0] != cgraph->nodes[norm_i - 1]) { + return match; + } + if (ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {}) || + ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {}) || + ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {}) || + ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_SCALE }, {})) { + return match; + } + const bool weighted = ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {}); + ggml_tensor * mul = weighted ? cgraph->nodes[norm_i + 1] : nullptr; + const bool closed = reshape + ? (weighted ? ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RESHAPE, GGML_OP_RMS_NORM, GGML_OP_MUL }, { i, i + 1, i + 3 }) : ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RESHAPE, GGML_OP_RMS_NORM }, { i, i + 1, i + 2 })) + : (weighted ? ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RMS_NORM, GGML_OP_MUL }, { i, i + 2 }) : ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RMS_NORM }, { i, i + 1 })); + if (closed && ggml_cuda_should_fuse_hc_post_norm(post, norm, mul)) { + match.post = post; + match.norm = norm; + match.mul = mul; + match.count = norm_i - i + 1 + int(weighted); + } + return match; +} + +static bool ggml_cuda_hc_post_norm_memory_ok(const ggml_cgraph * cgraph, int i, const ggml_cuda_hc_post_norm_match & match, int device) { + const auto range = [device](const ggml_tensor * tensor, uintptr_t & begin, uintptr_t & end) { + if (!tensor || !tensor->buffer || !tensor->data || !ggml_backend_buffer_is_cuda(tensor->buffer)) { + return false; + } + const auto * ctx = (const ggml_backend_cuda_buffer_context *) tensor->buffer->context; + if (ctx->device != device) { + return false; + } + begin = (uintptr_t) tensor->data; + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + const size_t alloc = ggml_backend_buffer_get_alloc_size(tensor->buffer, tensor); + if (begin < base || begin - base > size || alloc < ggml_nbytes(tensor) || alloc > size - (begin - base) || begin > UINTPTR_MAX - alloc || begin % sizeof(float) != 0) { + return false; + } + end = begin + alloc; + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + if (root->view_src || root->buffer != tensor->buffer || !root->data || (uintptr_t) root->data < base || + (uintptr_t) root->data - base > size || ggml_nbytes(root) > size - ((uintptr_t) root->data - base) || tensor->view_offs > ggml_nbytes(root) || + begin != (uintptr_t) root->data + tensor->view_offs || ggml_nbytes(tensor) > ggml_nbytes(root) - tensor->view_offs) { + return false; + } + } + return true; + }; + const ggml_tensor * writes[] = { match.post, match.mul ? match.mul : match.norm }; + uintptr_t begin[2], end[2]; + for (int w = 0; w < 2; ++w) { + if (!range(writes[w], begin[w], end[w])) { + return false; + } + } + for (int j = i + 1; j < i + match.count; ++j) { + if (cgraph->nodes[j]->op == GGML_OP_RESHAPE) { + uintptr_t lo, hi; + const ggml_tensor * view = cgraph->nodes[j]; + if (!range(view, lo, hi) || view->buffer != match.post->buffer || lo != begin[0] || ggml_nbytes(view) != ggml_nbytes(match.post)) { + return false; + } + } + } + const auto overlaps = [](uintptr_t a, uintptr_t b, uintptr_t c, uintptr_t d) { return a < d && c < b; }; + if (overlaps(begin[0], end[0], begin[1], end[1]) && (begin[0] != begin[1] || end[0] != end[1])) { + return false; + } + for (int j = i; j < i + match.count; ++j) { + for (const ggml_tensor * src : cgraph->nodes[j]->src) { + if (!src) { + continue; + } + bool internal = false; + for (int k = i; k < j; ++k) { + internal |= src == cgraph->nodes[k]; + } + if (internal) { + continue; + } + uintptr_t lo, hi; + if (!range(src, lo, hi)) { + return false; + } + for (int w = 0; w < 2; ++w) { + if (!overlaps(begin[w], end[w], lo, hi)) { + continue; + } + if (src == match.post->src[1] && !match.post->src[3] && + ggml_are_same_shape(match.post, src) && ggml_are_same_stride(match.post, src) && begin[w] == lo && end[w] == hi) { + continue; + } + return false; + } + } + } + return true; +} + +// try and fuse nodes and return the number of nodes to skip +static bool ggml_cuda_hc_injection_bytes(const ggml_tensor * tensor, size_t & bytes) { + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0) { return false; } + } + const size_t block = ggml_blck_size(tensor->type); + bytes = ggml_type_size(tensor->type); + int first = 0; + if (block > 1) { + if (!tensor->nb[0] || uint64_t(tensor->ne[0]) > SIZE_MAX/tensor->nb[0] || tensor->ne[0] % block) { + return false; + } + bytes = size_t(tensor->ne[0])*tensor->nb[0]/block; + first = 1; + } + for (int d = first; d < GGML_MAX_DIMS; ++d) { + const uint64_t count = tensor->ne[d] - 1; + if (count && tensor->nb[d] > (SIZE_MAX - bytes)/count) { return false; } + bytes += size_t(count)*tensor->nb[d]; + } + return bytes && bytes == ggml_nbytes(tensor); +} + +static bool ggml_cuda_hc_injection_range(const ggml_tensor * tensor, int device, uintptr_t & begin, uintptr_t & end) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + size_t bytes; + if (!ggml_cuda_hc_injection_bytes(tensor, bytes)) { return false; } + begin = (uintptr_t) tensor->data; + if (begin % sizeof(float)) { return false; } + if (begin < base || begin - base > size || bytes > size - (begin - base) || begin > UINTPTR_MAX - bytes) { + return false; + } + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + const uintptr_t root_begin = (uintptr_t) root->data; + size_t root_bytes; + if (!ggml_cuda_hc_injection_bytes(root, root_bytes)) { return false; } + if (root_begin % sizeof(float) || root->view_src || root->buffer != tensor->buffer || !root->data || root_begin < base || root_begin - base > size || + root_bytes > size - (root_begin - base) || tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs || + root_begin > UINTPTR_MAX - tensor->view_offs || begin != root_begin + tensor->view_offs) { + return false; + } + } + end = begin + bytes; + return true; +} + + +struct ggml_cuda_hc_injection_match { + ggml_tensor * first = nullptr; + ggml_tensor * unary = nullptr; + ggml_tensor * last = nullptr; + ggml_tensor * post = nullptr; + ggml_cuda_hc_post_norm_match hc; + int count = 0; +}; + +static ggml_cuda_hc_injection_match ggml_cuda_match_hc_injection(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_cuda_hc_injection_match match; + ggml_tensor * first = graph->nodes[i]; + const bool before = first->op == GGML_OP_SCALE; + const int u = i + int(before); + if (u >= graph->n_nodes || graph->nodes[u]->op != GGML_OP_UNARY) { return match; } + ggml_tensor * unary = graph->nodes[u]; + const ggml_unary_op op = ggml_get_unary_op(unary); + if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SILU) { return match; } + const bool after = u + 1 < graph->n_nodes && graph->nodes[u + 1]->op == GGML_OP_SCALE; + if (!before && !after) { return match; } + const int weight_last = u + int(after); + const int post_i = weight_last + 1; + if (post_i >= graph->n_nodes) { return match; } + ggml_tensor * last = graph->nodes[weight_last]; + ggml_tensor * post = graph->nodes[post_i]; + const ggml_tensor * raw = first->src[0]; + if (!raw || post->op != GGML_OP_DSV4_HC_POST || post->src[2] != last || + (before && unary->src[0] != first) || (after && last->src[0] != unary)) { return match; } + for (const ggml_tensor * tensor : {raw, (const ggml_tensor *) first, (const ggml_tensor *) unary, (const ggml_tensor *) last}) { + size_t bytes; + if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) || !ggml_are_same_shape(raw, tensor) || + !ggml_cuda_hc_injection_bytes(tensor, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } + } + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * comb = post->src[3]; + if (!x || !residual || x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || + !ggml_is_contiguous(post) || !ggml_are_same_shape(post, residual) || x->ne[0] <= 0 || residual->ne[1] <= 0 || x->ne[1] <= 0 || + x->ne[2] != 1 || x->ne[3] != 1 || residual->ne[0] != x->ne[0] || residual->ne[2] != x->ne[1] || residual->ne[3] != 1 || + raw->ne[0] != residual->ne[1] || raw->ne[1] != x->ne[1] || raw->ne[2] != 1 || raw->ne[3] != 1 || + (comb && (comb->type != GGML_TYPE_F32 || comb->ne[0] != residual->ne[1] || comb->ne[1] != residual->ne[1] || comb->ne[2] != x->ne[1] || comb->ne[3] != 1))) { return {}; } + size_t bytes; + if (!ggml_cuda_hc_injection_bytes(post, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } + auto hc = ggml_cuda_match_hc_post_norm(graph, post_i); + if (hc.count && graph->nodes[post_i + hc.count - 1] != (hc.mul ? hc.mul : hc.norm)) { return {}; } + if (hc.count && hc.norm->ne[0] != x->ne[0]) { + if (comb && bytes/sizeof(float) <= 256) { return {}; } + if (hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] > WARP_SIZE/2 && hc.norm->ne[0] < WARP_SIZE) { return {}; } + const int block_size = hc.norm->ne[0] < 1024 ? 256 : 1024; + const int64_t nweights = (hc.norm->ne[0] + x->ne[0] - 2)/x->ne[0] + 1; + if ((hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] <= WARP_SIZE/2) || nweights >= block_size) { hc = {}; } + } + const int end = hc.count ? post_i + hc.count : post_i + 1; + const int64_t ncols = hc.count ? hc.norm->ne[0] : 256; + const size_t shared = ((hc.count ? 32 : 0) + (ncols + x->ne[0] - 2)/x->ne[0] + 1)*sizeof(float); + if (shared > ggml_cuda_info().devices[ctx.device].smpb) { return {}; } + std::vector ops; + std::vector outputs; + for (int j = i; j < end; ++j) { + ops.push_back(graph->nodes[j]->op); + if (j <= weight_last || j == post_i || j == end - 1) { outputs.push_back(j); } + } + if (!ggml_can_fuse_subgraph(graph, i, ops.size(), ops.data(), outputs.data(), outputs.size())) { return {}; } + const ggml_tensor * gamma = hc.mul ? hc.mul->src[hc.mul->src[0] == hc.norm ? 1 : 0] : nullptr; + const ggml_tensor * reads[] = {raw, x, residual, comb, gamma}; + for (const ggml_tensor * tensor : reads) { + if (!tensor) { continue; } + const ggml_tensor * root = tensor->view_src ? tensor->view_src : tensor; + size_t span; + if (!ggml_cuda_hc_injection_bytes(tensor, span) || span > INT64_MAX) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { if (tensor->nb[d] % sizeof(float)) { return {}; } } + for (int j = i; j < end; ++j) { if (graph->nodes[j] == root) { return {}; } } + } + if (allocated) { + if (hc.count && !ggml_cuda_hc_post_norm_memory_ok(graph, post_i, hc, ctx.device)) { return {}; } + std::vector writes = {first, unary, last, post}; + if (hc.count) { writes.push_back(hc.mul ? hc.mul : hc.norm); } + std::vector wb(writes.size()), we(writes.size()); + for (size_t j = 0; j < writes.size(); ++j) { + if (!ggml_cuda_hc_injection_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } + for (size_t k = 0; k < j; ++k) { + if (wb[j] < we[k] && wb[k] < we[j] && + !((j < 3 && k < 3 && wb[j] == wb[k] && we[j] == we[k]) || (j >= 3 && k >= 3 && wb[j] == wb[k] && we[j] == we[k]))) { return {}; } + } + } + for (const ggml_tensor * tensor : reads) { + if (!tensor) { continue; } + uintptr_t lo, hi; + if (!ggml_cuda_hc_injection_range(tensor, ctx.device, lo, hi)) { return {}; } + for (size_t j = 0; j < writes.size(); ++j) { + if (!(wb[j] < hi && lo < we[j])) { continue; } + if (j >= 3 && tensor == residual && !comb && ggml_are_same_shape(writes[j], residual) && + ggml_are_same_stride(writes[j], residual) && wb[j] == lo && we[j] == hi) { continue; } + return {}; + } + } + } + match.first = first; + match.unary = unary; + match.last = last; + match.post = post; + match.hc = hc; + match.count = end - i; + return match; +} + +static bool ggml_cuda_hc_up_bytes(const ggml_tensor * tensor, size_t & bytes) { + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0) { return false; } + } + const size_t block = ggml_blck_size(tensor->type); + bytes = ggml_type_size(tensor->type); + int first = 0; + if (block > 1) { + if (!tensor->nb[0] || uint64_t(tensor->ne[0]) > SIZE_MAX/tensor->nb[0] || tensor->ne[0] % block) { + return false; + } + bytes = size_t(tensor->ne[0])*tensor->nb[0]/block; + first = 1; + } + for (int d = first; d < GGML_MAX_DIMS; ++d) { + const uint64_t count = tensor->ne[d] - 1; + if (count && tensor->nb[d] > (SIZE_MAX - bytes)/count) { return false; } + bytes += size_t(count)*tensor->nb[d]; + } + return bytes && bytes == ggml_nbytes(tensor); +} + +static bool ggml_cuda_hc_up_range(const ggml_tensor * tensor, int device, uintptr_t & begin, uintptr_t & end) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + size_t bytes; + if (!ggml_cuda_hc_up_bytes(tensor, bytes)) { return false; } + begin = (uintptr_t) tensor->data; + if (begin < base || begin - base > size || bytes > size - (begin - base) || begin > UINTPTR_MAX - bytes) { + return false; + } + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + const uintptr_t root_begin = (uintptr_t) root->data; + size_t root_bytes; + if (!ggml_cuda_hc_up_bytes(root, root_bytes)) { return false; } + if (root->view_src || root->buffer != tensor->buffer || !root->data || root_begin < base || root_begin - base > size || + root_bytes > size - (root_begin - base) || tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs || + root_begin > UINTPTR_MAX - tensor->view_offs || begin != root_begin + tensor->view_offs) { + return false; + } + } + end = begin + bytes; + return true; +} + +struct ggml_cuda_hc_up_match { + ggml_tensor * mm = nullptr; + ggml_tensor * dst = nullptr; + int count = 0; +}; + +static ggml_cuda_hc_up_match ggml_cuda_match_hc_up_shape(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i) { + ggml_tensor * mm = graph->nodes[i]; + if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || + mm->src[1]->type != GGML_TYPE_F32 || mm->ne[2] != 1 || mm->ne[3] != 1 || mm->ne[1] > MMVF_MAX_BATCH_SIZE || + ggml_cuda_op_mul_mat_use_fwht(mm)) { + return {}; + } + int last = i + 1; + while (last < graph->n_nodes && last < i + 3 && graph->nodes[last]->op == GGML_OP_RESHAPE) { ++last; } + if (last >= graph->n_nodes || graph->nodes[last]->op != GGML_OP_DSV4_HC_PRE) { return {}; } + ggml_tensor * dst = graph->nodes[last]; + const ggml_tensor * norm = dst->src[0]; + const ggml_tensor * view = dst->src[1]; + const ggml_tensor * weight = mm->src[0]; + const ggml_tensor * input = mm->src[1]; + if (!norm || !view || !ggml_get_op_params_i32(dst, 1) || norm->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + norm->ne[0] <= 0 || norm->ne[1] <= 0 || norm->ne[2] <= 0 || view->op != GGML_OP_RESHAPE || view->src[0] != mm || view->type != GGML_TYPE_F32 || !ggml_is_contiguous(view) || + !ggml_are_same_shape(view, norm) || norm->ne[3] != 1 || mm->ne[0]/norm->ne[0] != norm->ne[1] || mm->ne[0] % norm->ne[0] || mm->ne[1] != norm->ne[2] || + dst->ne[0] != norm->ne[0] || dst->ne[1] != norm->ne[2] || dst->ne[2] != 1 || dst->ne[3] != 1 || + weight->ne[2] != 1 || weight->ne[3] != 1 || input->ne[2] != 1 || input->ne[3] != 1 || + (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16 && weight->type != GGML_TYPE_BF16)) { + return {}; + } + if (!ggml_cuda_should_use_mmvf(weight->type, ggml_cuda_info().devices[ctx.device].cc, weight->ne, weight->nb, input->ne[1]) || + !ggml_cuda_should_fuse_hc_up(weight, norm, ctx.device)) { return {}; } + const int count = last - i + 1; + ggml_op ops[4]; int outputs[4]; + for (int j = 0; j < count; ++j) { ops[j] = graph->nodes[i + j]->op; outputs[j] = i + j; } + if (!ggml_can_fuse_subgraph(graph, i, count, ops, outputs, count)) { return {}; } + return {mm, dst, count}; +} + +static ggml_cuda_hc_up_match ggml_cuda_match_hc_up(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i) { + const auto match = ggml_cuda_match_hc_up_shape(ctx, graph, i); + if (!match.count) { return {}; } + ggml_tensor * mm = match.mm; + ggml_tensor * dst = match.dst; + const ggml_tensor * norm = dst->src[0]; + const ggml_tensor * view = dst->src[1]; + const ggml_tensor * weight = mm->src[0]; + const ggml_tensor * input = mm->src[1]; + const ggml_tensor * tensors[] = {weight, input, mm}; + for (const ggml_tensor * tensor : tensors) { + uintptr_t begin, end; + if (!ggml_cuda_hc_up_range(tensor, ctx.device, begin, end) || tensor->nb[0] != ggml_type_size(tensor->type) || + (end - begin)/ggml_type_size(tensor->type) > INT_MAX || begin % (2*ggml_type_size(tensor->type))) { return {}; } + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] > INT_MAX || tensor->nb[d]/ggml_type_size(tensor->type) > INT_MAX || + (tensor != mm && tensor->nb[d] % (2*ggml_type_size(tensor->type)))) { return {}; } + } + } + if (!(mm->flags & GGML_TENSOR_FLAG_COMPUTE) || !(dst->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_is_contiguous(mm) || + !ggml_is_contiguous(dst) || (ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + weight->view_src && ggml_nbytes(weight) != ggml_backend_buffer_get_alloc_size(weight->buffer, weight))) { return {}; } + uintptr_t begin[2], end[2], reads_begin[3], reads_end[3], view_begin, view_end; + const ggml_tensor * reads[] = {weight, input, norm}; + for (int j = 0; j < 3; ++j) { + if (!ggml_cuda_hc_up_range(reads[j], ctx.device, reads_begin[j], reads_end[j]) || reads_begin[j] % ggml_type_size(reads[j]->type)) { return {}; } + } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (norm->nb[d] % sizeof(float) || norm->nb[d]/sizeof(float) > INT64_MAX) { return {}; } + } + if (!ggml_cuda_hc_up_range(mm, ctx.device, begin[0], end[0]) || !ggml_cuda_hc_up_range(dst, ctx.device, begin[1], end[1]) || + !ggml_cuda_hc_up_range(view, ctx.device, view_begin, view_end) || view_begin != begin[0] || view_end != end[0]) { return {}; } + for (int j = 0; j < 2; ++j) { + if (begin[j] % sizeof(float)) { return {}; } + for (int k = 0; k < 3; ++k) { + if (begin[j] < reads_end[k] && reads_begin[k] < end[j]) { return {}; } + } + } + if (begin[0] < end[1] && begin[1] < end[0]) { return {}; } + return match; +} + +static bool ggml_cuda_affine_unary_bytes(const ggml_tensor * tensor, size_t & bytes) { + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0) { return false; } + } + const size_t block = ggml_blck_size(tensor->type); + bytes = ggml_type_size(tensor->type); + int first = 0; + if (block > 1) { + if (!tensor->nb[0] || uint64_t(tensor->ne[0]) > SIZE_MAX/tensor->nb[0] || tensor->ne[0] % block) { + return false; + } + bytes = size_t(tensor->ne[0])*tensor->nb[0]/block; + first = 1; + } + for (int d = first; d < GGML_MAX_DIMS; ++d) { + const uint64_t count = tensor->ne[d] - 1; + if (count && tensor->nb[d] > (SIZE_MAX - bytes)/count) { return false; } + bytes += size_t(count)*tensor->nb[d]; + } + return bytes && bytes == ggml_nbytes(tensor); +} + +static bool ggml_cuda_affine_unary_range(const ggml_tensor * tensor, int device, uintptr_t & begin, uintptr_t & end) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(tensor, bytes)) { return false; } + begin = (uintptr_t) tensor->data; + if (begin < base || begin - base > size || bytes > size - (begin - base) || begin > UINTPTR_MAX - bytes) { + return false; + } + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + const uintptr_t root_begin = (uintptr_t) root->data; + size_t root_bytes; + if (!ggml_cuda_affine_unary_bytes(root, root_bytes)) { return false; } + if (root->view_src || root->buffer != tensor->buffer || !root->data || root_begin < base || root_begin - base > size || + root_bytes > size - (root_begin - base) || tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs || + root_begin > UINTPTR_MAX - tensor->view_offs || begin != root_begin + tensor->view_offs) { + return false; + } + } + end = begin + bytes; + return true; +} + +struct ggml_cuda_affine_unary_match { + ggml_tensor * mul = nullptr; + ggml_tensor * add = nullptr; + ggml_tensor * unary = nullptr; + ggml_tensor * post = nullptr; + int count = 0; +}; + +static bool ggml_cuda_affine_unary_binary_type(const ggml_tensor * node) { + if (!node->src[0] || !node->src[1]) { return false; } + const ggml_type a = node->src[0]->type; + const ggml_type b = node->src[1]->type; + return (a == GGML_TYPE_F32 && b == GGML_TYPE_F32 && node->type == GGML_TYPE_F32) || + (a == GGML_TYPE_F16 && b == GGML_TYPE_F16 && node->type == GGML_TYPE_F16) || + (a == GGML_TYPE_F16 && b == GGML_TYPE_F32 && (node->type == GGML_TYPE_F16 || node->type == GGML_TYPE_F32)) || + (a == GGML_TYPE_BF16 && (b == GGML_TYPE_F32 || b == GGML_TYPE_BF16) && node->type == GGML_TYPE_BF16); +} + +static bool ggml_cuda_affine_unary_prior_view(ggml_cgraph * graph, int start, const ggml_tensor * tensor) { + const ggml_tensor * root = tensor->view_src; + size_t bytes, root_bytes; + if (tensor->op != GGML_OP_VIEW || !root || root->view_src || + !ggml_cuda_affine_unary_bytes(tensor, bytes) || !ggml_cuda_affine_unary_bytes(root, root_bytes) || + tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs) { + return false; + } + if (root->op == GGML_OP_NONE) { return true; } + for (int i = 0; i < start; ++i) { + if (graph->nodes[i] == root) { return true; } + } + return false; +} + +static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_cuda_affine_unary_match match; + ggml_tensor * mul = graph->nodes[i]; + if (mul->op != GGML_OP_MUL || !ggml_cuda_affine_unary_binary_type(mul)) { return match; } + const auto next_compute = [&](int index) { + while (index < graph->n_nodes && graph->nodes[index]->op == GGML_OP_VIEW) { + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[index])) { return graph->n_nodes; } + uintptr_t begin, end; + if (allocated && !ggml_cuda_affine_unary_range(graph->nodes[index], ctx.device, begin, end)) { return graph->n_nodes; } + ++index; + } + return index; + }; + const int a = next_compute(i + 1); + if (a >= graph->n_nodes || graph->nodes[a]->op != GGML_OP_ADD || graph->nodes[a]->src[0] != mul || + !ggml_cuda_affine_unary_binary_type(graph->nodes[a])) { return match; } + ggml_tensor * add = graph->nodes[a]; + const int u = next_compute(a + 1); + if (u >= graph->n_nodes || graph->nodes[u]->op != GGML_OP_UNARY || graph->nodes[u]->src[0] != add || + graph->nodes[u]->type != add->type) { return match; } + ggml_tensor * unary = graph->nodes[u]; + const ggml_unary_op op = ggml_get_unary_op(unary); + if ((op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SILU) || + ggml_cuda_can_fuse(graph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { op })) { return match; } + ggml_cuda_topk_moe_args topk{}; + if (op == GGML_UNARY_OP_SIGMOID && ggml_cuda_topk_moe_fusion(graph, u, topk)) { return match; } + const int p = next_compute(u + 1); + ggml_tensor * post = p < graph->n_nodes && graph->nodes[p]->op == GGML_OP_SCALE && graph->nodes[p]->src[0] == unary ? graph->nodes[p] : nullptr; + if (post && ((post->type != GGML_TYPE_F32 && post->type != GGML_TYPE_BF16) || post->type != unary->type)) { return match; } + const int indices[] = { i, a, u, p }; + const ggml_op ops[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY, GGML_OP_SCALE }; + const int n = post ? 4 : 3; + if (!ggml_can_fuse_subgraph_ext(graph, indices, n, ops, indices, n)) { return match; } + int64_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (mul->ne[d] <= 0 || mul->ne[d] > (INT_MAX - CUDA_NEG_BLOCK_SIZE)/elements) { return {}; } + elements *= mul->ne[d]; + } + const ggml_tensor * reads[] = { mul->src[0], mul->src[1], add->src[1] }; + uintptr_t read_begin[3], read_end[3], begin[4], end[4]; + for (int j = 0; j < 3; ++j) { + const ggml_tensor * tensor = reads[j]; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(tensor, bytes) || + (allocated && (!ggml_cuda_affine_unary_range(tensor, ctx.device, read_begin[j], read_end[j]) || + read_begin[j] % ggml_type_size(tensor->type)))) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] > INT_MAX || tensor->nb[d] % ggml_type_size(tensor->type) || + tensor->nb[d]/ggml_type_size(tensor->type) > UINT32_MAX || + mul->ne[d] % tensor->ne[d]) { return {}; } + } + for (int k = 0; k < n; ++k) { + if (tensor == graph->nodes[indices[k]] || tensor->view_src == graph->nodes[indices[k]]) { return {}; } + } + } + if (!ggml_are_same_shape(reads[0], mul)) { return {}; } + for (int j = 0; j < n; ++j) { + const ggml_tensor * tensor = graph->nodes[indices[j]]; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(tensor, bytes) || !ggml_is_contiguous(tensor) || !ggml_are_same_shape(mul, tensor) || + (allocated && (!ggml_cuda_affine_unary_range(tensor, ctx.device, begin[j], end[j]) || + begin[j] % ggml_type_size(tensor->type)))) { return {}; } + if (!allocated) { continue; } + for (int k = 0; k < 3; ++k) { + if (begin[j] < read_end[k] && read_begin[k] < end[j] && + (k != 0 || begin[j] != read_begin[k] || end[j] != read_end[k] || tensor->type != reads[0]->type || !ggml_is_contiguous(reads[0]))) { + return {}; + } + } + for (int k = 0; k < j; ++k) { + if (begin[j] < end[k] && begin[k] < end[j] && + (begin[j] != begin[k] || end[j] != end[k] || tensor->type != graph->nodes[indices[k]]->type)) { return {}; } + } + } + match.mul = mul; + match.add = add; + match.unary = unary; + match.post = post; + match.count = indices[n - 1] - i + 1; + return match; +} + // 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) { @@ -4150,6 +4724,22 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return fused_node_count - 1; } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto injection = ggml_cuda_match_hc_injection(*cuda_ctx, cgraph, i); + if (injection.count) { + ggml_cuda_op_hc_injection(*cuda_ctx, injection.first, injection.unary, injection.last, injection.post, injection.hc.norm, injection.hc.mul); + return injection.count - 1; + } + } + + if (node->op == GGML_OP_DSV4_HC_POST && cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_hc_post_norm(cgraph, i); + if (match.count != 0 && ggml_cuda_hc_post_norm_memory_ok(cgraph, i, match, cuda_ctx->device)) { + ggml_cuda_op_hc_post_norm(*cuda_ctx, match.post, match.norm, match.mul); + return match.count - 1; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]); return 4; @@ -4202,6 +4792,22 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); + if (affine.count) { + ggml_cuda_op_affine_unary(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post); + return affine.count - 1; + } + } + + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_hc_up(*cuda_ctx, cgraph, i); + if (match.count) { + ggml_cuda_mul_mat_vec_f_hc_pre(*cuda_ctx, match.mm, match.dst); + return match.count - 1; + } + } + return 0; } @@ -4546,6 +5152,55 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { + const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); + if (affine.count) { + ggml_tensor * last = cgraph->nodes[i + affine.count - 1]; + params->add_alloc_dep(params->user_data, affine.mul->src[0], last); + params->add_alloc_dep(params->user_data, affine.mul->src[1], last); + params->add_alloc_dep(params->user_data, affine.add->src[1], last); + i += affine.count - 1; + continue; + } + const auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); + if (hc_up.count) { + params->add_alloc_dep(params->user_data, hc_up.mm->src[0], hc_up.dst); + params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); + } + + const auto injection = ggml_cuda_match_hc_injection(*cuda_ctx, cgraph, i, false); + if (injection.count) { + ggml_tensor * terminal = cgraph->nodes[i + injection.count - 1]; + params->add_alloc_dep(params->user_data, injection.first->src[0], terminal); + for (ggml_tensor * tensor : {injection.first, injection.unary, injection.last}) { + if (tensor != terminal) { params->add_alloc_dep(params->user_data, tensor, terminal); } + } + for (ggml_tensor * tensor : {injection.post->src[0], injection.post->src[1], injection.post->src[3]}) { + if (tensor) { params->add_alloc_dep(params->user_data, tensor, terminal); } + } + if (injection.hc.mul) { + params->add_alloc_dep(params->user_data, injection.hc.mul->src[injection.hc.mul->src[0] == injection.hc.norm ? 1 : 0], terminal); + } + i += injection.count - 1; + continue; + } + if (cgraph->nodes[i]->op == GGML_OP_DSV4_HC_POST) { + const auto match = ggml_cuda_match_hc_post_norm(cgraph, i); + if (match.count != 0) { + for (int j = i; j < i + match.count; ++j) { + for (const ggml_tensor * src : cgraph->nodes[j]->src) { + bool internal = false; + for (int k = i; src && k < j; ++k) { + internal |= src == cgraph->nodes[k]; + } + if (src && !internal) { + params->add_alloc_dep(params->user_data, const_cast(src), cgraph->nodes[i + match.count - 1]); + } + } + } + i += match.count - 1; + continue; + } + } ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { params->add_alloc_dep(params->user_data, const_cast(match.experts), match.dst); diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index bd5c5d421a4..44820f029ca 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -4,18 +4,40 @@ #include "mmvf.cuh" #include "convert.cuh" -template -static __global__ void mul_mat_vec_f( +struct ggml_cuda_mmvf_store { + __device__ __forceinline__ void operator()(float * dst, const float * base, int index, float value) const { + dst[index] = value; + GGML_UNUSED(base); + } +}; + +struct ggml_cuda_mmvf_geometry { + __device__ __forceinline__ int row() const { return blockIdx.x; } + __device__ __forceinline__ int shared_offset() const { return 0; } + __device__ __forceinline__ void sync() const { ggml_cuda_pdl_sync(); } + __device__ __forceinline__ void launch_complete() const { ggml_cuda_pdl_lc(); } +}; + +struct ggml_cuda_mmvf_hc_geometry { + int index; + __device__ __forceinline__ int row() const { return index; } + __device__ __forceinline__ int shared_offset() const { return threadIdx.y*ggml_cuda_get_physical_warp_size()*sizeof(float); } + __device__ __forceinline__ void sync() const {} + __device__ __forceinline__ void launch_complete() const {} +}; + +template +static __device__ __forceinline__ void mul_mat_vec_f_impl( const T * x_ptr, const float * y_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const int ncols2, const uint3 nchannels_y, const int stride_row, const int stride_col_y2, const int stride_col_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const uint3 sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride) { + const int ids_stride, const Write write, const Geometry geometry = {}) { const T * GGML_CUDA_RESTRICT x = x_ptr; const float * GGML_CUDA_RESTRICT y = y_ptr; const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; - const int row = blockIdx.x; + const int row = geometry.row(); // for MUL_MAT_ID - blockIdx.y = n_expert_used, blockIdx.z = ncols_dst (tokens) const int channel_dst = blockIdx.y; const int tid = threadIdx.x; @@ -25,7 +47,7 @@ static __global__ void mul_mat_vec_f( int channel_y; int sample_dst; - ggml_cuda_pdl_sync(); + geometry.sync(); if constexpr (is_multi_token_id) { // Multi-token MUL_MAT_ID path, adding these in the normal path causes a perf regression for n_tokens=1 case token_idx = blockIdx.z; @@ -99,7 +121,7 @@ static __global__ void mul_mat_vec_f( const float2 * y2 = (const float2 *) y; extern __shared__ char data_mmv[]; - float * buf_iw = (float *) data_mmv; + float * buf_iw = (float *) (data_mmv + geometry.shared_offset()); [[maybe_unused]] float * buf_iw_gate = nullptr; if constexpr (has_fusion) { buf_iw_gate = (float *) (data_mmv + warp_size*sizeof(float)); @@ -304,7 +326,7 @@ static __global__ void mul_mat_vec_f( static_assert(std::is_same_v, "unsupported type"); } - ggml_cuda_pdl_lc(); + geometry.launch_complete(); #pragma unroll for (int j = 0; j < ncols_dst; ++j) { sumf[j] = warp_reduce_sum(sumf[j]); @@ -376,13 +398,24 @@ static __global__ void mul_mat_vec_f( } } - dst[tid*stride_col_dst + row] = value; + write(dst, dst_ptr, tid*stride_col_dst + row, value); if constexpr (!has_fusion) { GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, glu_limit, gate_x, x_bias, gate_bias, sumf_gate); } } +template +static __global__ void mul_mat_vec_f( + const T * x_ptr, const float * y_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, + const int ncols2, const uint3 nchannels_y, const int stride_row, const int stride_col_y2, const int stride_col_dst, + const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const uint3 sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const int ids_stride) { + mul_mat_vec_f_impl( + x_ptr, y_ptr, ids_ptr, fusion, dst_ptr, ncols2, nchannels_y, stride_row, stride_col_y2, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ggml_cuda_mmvf_store{}); +} + template static void mul_mat_vec_f_switch_fusion( const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -873,3 +906,139 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 return false; } } + +struct ggml_cuda_mmvf_hc_store { + const float * norm; + float * sum; + int i0; + int64_t ih; + int64_t stride0; + int64_t stride1; + int64_t stride2; + __device__ __forceinline__ void operator()(float * dst, const float * base, int index, float value) const { + dst[index] = value; + const float xv = norm[int64_t(i0)*stride0 + ih*stride1 + int64_t(threadIdx.x)*stride2]; + const float wv = 1.0f / (1.0f + expf(-value)); + *sum += xv * wv; + GGML_UNUSED(base); + } +}; + +template +static __global__ void mul_mat_vec_f_hc_pre( + const T * weight, const float * input, const float * norm, float * gate, float * dst, + int k2, int weight_stride, int input_stride, int gate_stride, int n, int hc, + int64_t norm_stride0, int64_t norm_stride1, int64_t norm_stride2, int64_t dst_stride, float scale) { + const uint3 one = make_uint3(1, 0, 1); + const int i0 = blockIdx.x; + const int tid = threadIdx.x; + float sum = 0.0f; + ggml_cuda_pdl_sync(); + for (int64_t ih = 0; ih < hc; ++ih) { + const int row = i0 + ih*n; + const ggml_cuda_mmvf_hc_store write = {norm, &sum, i0, ih, norm_stride0, norm_stride1, norm_stride2}; + const ggml_cuda_mmvf_hc_geometry geometry = {row}; + mul_mat_vec_f_impl( + weight, input, nullptr, {}, gate, k2, make_uint3(0, 0, 0), weight_stride, input_stride, gate_stride, + one, 0, 0, 0, one, 0, 0, 0, 0, write, geometry); + } + ggml_cuda_pdl_lc(); + if (tid < columns) { + dst[int64_t(tid)*dst_stride + i0] = scale * sum; + } +} + +static int ggml_cuda_mmvf_hc_block_size(int64_t ncols, int device) { + const int warp_size = ggml_cuda_info().devices[device].warp_size; + int64_t best = warp_size; + int64_t iterations = (ncols + 2*warp_size - 1)/(2*warp_size); + int64_t maximum = 256; + const int cc = ggml_cuda_info().devices[device].cc; + if (cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_RDNA1) { maximum = 128; } + for (int64_t block = 2*warp_size; block <= maximum; block += warp_size) { + const int64_t count = (ncols + 2*block - 1)/(2*block); + if (count < iterations) { iterations = count; best = block; } + } + return best; +} + +bool ggml_cuda_should_fuse_hc_up(const ggml_tensor * weight, const ggml_tensor * norm, int device) { + static const auto cache_sizes = []() { + std::array sizes = {}; + const auto & info = ggml_cuda_info(); + for (int i = 0; i < info.device_count; ++i) { + cudaDeviceProp props = {}; + if (cudaGetDeviceProperties(&props, info.devices[i].physical_device) == cudaSuccess) { + sizes[i] = props.l2CacheSize; + } + } + return sizes; + }(); + const auto & info = ggml_cuda_info().devices[device]; + // Keep enough CTAs and leave cache space for the preceding projection. + if (norm->ne[0] < int64_t(4)*info.nsm || weight->ne[1] <= 0 || weight->nb[1] == 0 || cache_sizes[device] <= 0) { + return false; + } + return weight->nb[1] <= size_t(cache_sizes[device])/2/uint64_t(weight->ne[1]); +} + +template +static void ggml_cuda_mmvf_hc_launch(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, const ggml_tensor * dst) { + const ggml_tensor * weight = mm->src[0]; + const ggml_tensor * input = mm->src[1]; + const ggml_tensor * norm = dst->src[0]; + auto kernel = mul_mat_vec_f_hc_pre; + const int shared = ggml_cuda_info().devices[ctx.device].warp_size*sizeof(float); + const ggml_cuda_kernel_launch_params params = {dim3(norm->ne[0], 1, 1), dim3(Block, 1, 1), shared, ctx.stream()}; + ggml_cuda_kernel_launch(kernel, params, + (const T *) weight->data, (const float *) input->data, (const float *) norm->data, (float *) mm->data, (float *) dst->data, + int(weight->ne[0]/2), int(weight->nb[1]/sizeof(T)), int(input->nb[1]/(2*sizeof(float))), int(mm->nb[1]/sizeof(float)), + int(norm->ne[0]), int(norm->ne[1]), int64_t(norm->nb[0]/sizeof(float)), int64_t(norm->nb[1]/sizeof(float)), + int64_t(norm->nb[2]/sizeof(float)), int64_t(dst->nb[1]/sizeof(float)), ggml_get_op_params_f32(dst, 0)); +} + +template +static void ggml_cuda_mmvf_hc_blocks(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, const ggml_tensor * dst) { + switch (ggml_cuda_mmvf_hc_block_size(mm->src[0]->ne[0], ctx.device)) { + case 32: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 64: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 96: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 128: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 160: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 192: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 224: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + case 256: ggml_cuda_mmvf_hc_launch(ctx, mm, dst); break; + default: GGML_ABORT("unsupported MMVF block size"); + } +} + +template +static void ggml_cuda_mmvf_hc_columns(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, const ggml_tensor * dst) { + switch (mm->ne[1]) { + case 1: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 2: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 3: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 4: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 5: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 6: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 7: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + case 8: ggml_cuda_mmvf_hc_blocks(ctx, mm, dst); break; + default: GGML_ABORT("unsupported MMVF column count"); + } +} + +void ggml_cuda_mul_mat_vec_f_hc_pre(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, const ggml_tensor * dst) { + const auto type = mm->src[0]->type; + const auto prec = fast_fp16_available(ggml_cuda_info().devices[ctx.device].cc) ? ggml_prec(mm->op_params[0]) : GGML_PREC_F32; + if (type == GGML_TYPE_F32) { + ggml_cuda_mmvf_hc_columns(ctx, mm, dst); + } else if (type == GGML_TYPE_F16 && prec == GGML_PREC_DEFAULT) { + ggml_cuda_mmvf_hc_columns(ctx, mm, dst); + } else if (type == GGML_TYPE_F16) { + ggml_cuda_mmvf_hc_columns(ctx, mm, dst); + } else if (type == GGML_TYPE_BF16) { + ggml_cuda_mmvf_hc_columns(ctx, mm, dst); + } else { + GGML_ABORT("unsupported MMVF type"); + } +} diff --git a/ggml/src/ggml-cuda/mmvf.cuh b/ggml/src/ggml-cuda/mmvf.cuh index a50f7c02180..5a5ee48eb25 100644 --- a/ggml/src/ggml-cuda/mmvf.cuh +++ b/ggml/src/ggml-cuda/mmvf.cuh @@ -12,3 +12,7 @@ void ggml_cuda_op_mul_mat_vec_f( const int64_t src1_padded_row_size, cudaStream_t stream); bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11); + +void ggml_cuda_mul_mat_vec_f_hc_pre(ggml_backend_cuda_context & ctx, const ggml_tensor * mm, const ggml_tensor * dst); + +bool ggml_cuda_should_fuse_hc_up(const ggml_tensor * weight, const ggml_tensor * norm, int device); diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 5543307be69..e5b6f2b18f8 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1,4 +1,5 @@ #include "norm.cuh" +#include "unary.cuh" #include template @@ -729,3 +730,342 @@ void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { l2_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } + +struct hc_post_norm_data { + const float * x; + const float * residual; + const float * post; + const float * comb; + const float * mul; + float * post_dst; + float * dst; + int64_t n_embd; + int64_t hc; + int ncols; + int64_t sx[2]; + int64_t sr[3]; + int64_t sp[2]; + int64_t sc[3]; + int64_t sm[4]; + uint3 mul_ne[4]; + uint3 embd_ne; + float eps; +}; + +template +static __global__ __launch_bounds__(block_size) void hc_post_norm_f32(hc_post_norm_data a) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; + } + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + ggml_cuda_pdl_sync(); + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * a.post[ih*a.sp[0] + it*a.sp[1]]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + a.dst[offset + col] = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + } else { + a.dst[offset + col] = scale * a.post_dst[offset + col]; + } + } +} + +template +static auto hc_post_norm_kernel(int layout) { + if (layout == 1) { + return hc_post_norm_f32; + } + if (layout == 2) { + return hc_post_norm_f32; + } + return hc_post_norm_f32; +} + +bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * mul) { + const ggml_tensor * dst = mul ? mul : norm; + if (!post->src[0] || !post->src[1] || !post->src[2] || post->op != GGML_OP_DSV4_HC_POST || norm->op != GGML_OP_RMS_NORM || + post->type != GGML_TYPE_F32 || norm->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + !ggml_is_contiguous(post) || !ggml_is_contiguous(norm) || !ggml_is_contiguous(dst) || + norm->ne[0] > INT_MAX - 1024 || norm->ne[1] > INT_MAX || norm->ne[2] > 65535 || norm->ne[3] > 65535) { + return false; + } + const ggml_tensor * inputs[] = { post->src[0], post->src[1], post->src[2], post->src[3], post, norm, dst, mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr }; + for (const ggml_tensor * tensor : inputs) { + if (!tensor) { + continue; + } + if (tensor->type != GGML_TYPE_F32) { + return false; + } + size_t span = sizeof(float); + size_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0 || tensor->ne[d] > INT_MAX || tensor->nb[d] % sizeof(float) != 0 || + elements > INT64_MAX/sizeof(float)/(size_t) tensor->ne[d] || + (tensor->nb[d] != 0 && (size_t) (tensor->ne[d] - 1) > (SIZE_MAX - span)/tensor->nb[d])) { + return false; + } + elements *= tensor->ne[d]; + span += (tensor->ne[d] - 1)*tensor->nb[d]; + } + if (span > INT64_MAX) { + return false; + } + } + const float eps = ggml_get_op_params_f32(norm, 0); + return ggml_nelements(norm) <= INT_MAX && ggml_nelements(post) == ggml_nelements(norm) && ggml_are_same_shape(norm, dst) && std::isfinite(eps) && eps >= 0.0f; +} + +void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = (float *) (mul ? mul : norm)->data; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm->ne[0]; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = ggml_get_op_params_f32(norm, 0); + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_post_norm_kernel<256, true, true>(layout) : hc_post_norm_kernel<256, true, false>(layout)) : (mul ? hc_post_norm_kernel<256, false, true>(layout) : hc_post_norm_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_post_norm_kernel<1024, true, true>(layout) : hc_post_norm_kernel<1024, true, false>(layout)) : (mul ? hc_post_norm_kernel<1024, false, true>(layout) : hc_post_norm_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a); +} + +struct hc_injection_data { + const float * input; + float * pre; + float * unary; + float * last; + float scale0; + float bias0; + float scale1; + float bias1; + ggml_unary_op op; + int64_t values; +}; + +static __device__ __forceinline__ float hc_injection_gate(hc_injection_data g, int64_t index, bool store) { + float value = g.input[index]; + if (g.pre) { + value = __fmaf_rn(g.scale0, value, g.bias0); + if (store) { g.pre[index] = value; } + } + value = g.op == GGML_UNARY_OP_SILU ? ggml_cuda_op_silu_single(value) : 1.0f / (1.0f + expf(-value)); + if (store) { g.unary[index] = value; } + if (g.last) { + value = __fmaf_rn(g.scale1, value, g.bias1); + if (store) { g.last[index] = value; } + } + return value; +} + +static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { + for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { + const int64_t scalar = first + j; + weights[j] = hc_injection_gate(g, scalar, scalar*n_embd >= begin); + } + __syncthreads(); +} + +template +static __global__ __launch_bounds__(block_size) void hc_injection_norm_f32(hc_post_norm_data a, hc_injection_data g) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + ggml_cuda_pdl_sync(); + extern __shared__ float s_sum[]; + float * weights = s_sum + 32; + const int64_t gate_first = offset/a.n_embd; + float grouped_weight = 0.0f; + if constexpr (grouped) { + if (threadIdx.x % WARP_SIZE == 0) { + grouped_weight = hc_injection_gate(g, gate_first, threadIdx.x == 0); + } + grouped_weight = __shfl_sync(0xffffffff, grouped_weight, 0, WARP_SIZE); + } else { + const int64_t gate_count = (offset + a.ncols - 1)/a.n_embd - gate_first + 1; + hc_injection_load(g, weights, gate_first, gate_count, offset, a.n_embd); + } + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; + } + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + const float weight = grouped ? grouped_weight : weights[ih + it*a.hc - gate_first]; + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weight; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + a.dst[offset + col] = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + } else { + a.dst[offset + col] = scale * a.post_dst[offset + col]; + } + } +} + +template +static __global__ void hc_injection_post_f32(hc_post_norm_data a, hc_injection_data g) { + ggml_cuda_pdl_lc(); + const int64_t begin = int64_t(blockIdx.x)*blockDim.x; + const int64_t end = begin + int64_t(blockDim.x) < g.values ? begin + int64_t(blockDim.x) : g.values; + const int64_t ir = begin + threadIdx.x; + ggml_cuda_pdl_sync(); + extern __shared__ float weights[]; + const int64_t first = begin/a.n_embd; + const int64_t count = (end - 1)/a.n_embd - first + 1; + hc_injection_load(g, weights, first, count, begin, a.n_embd); + if (ir >= g.values) { return; } + const int64_t i0 = ir % a.n_embd; + const int64_t ih = (ir/a.n_embd) % a.hc; + const int64_t it = ir/(a.n_embd*a.hc); + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weights[ir/a.n_embd - first]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; +} + +template +static auto hc_injection_norm_kernel(int layout) { + if (layout == 1) { return hc_injection_norm_f32; } + if (layout == 2) { return hc_injection_norm_f32; } + return hc_injection_norm_f32; +} + +void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = norm ? (float *) (mul ? mul : norm)->data : nullptr; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm ? norm->ne[0] : 256; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = norm ? ggml_get_op_params_f32(norm, 0) : 0.0f; + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + hc_injection_data g{}; + g.input = (const float *) first->src[0]->data; + g.pre = first != unary ? (float *) first->data : nullptr; + g.unary = (float *) unary->data; + g.last = last != unary ? (float *) last->data : nullptr; + g.scale0 = g.pre ? ggml_get_op_params_f32(first, 0) : 1.0f; + g.bias0 = g.pre ? ggml_get_op_params_f32(first, 1) : 0.0f; + g.scale1 = g.last ? ggml_get_op_params_f32(last, 0) : 1.0f; + g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; + g.op = ggml_get_unary_op(unary); + g.values = ggml_nelements(post); + if (!norm) { + const size_t shared = ((256 + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const ggml_cuda_kernel_launch_params launch(dim3((g.values + 255)/256), dim3(256), shared, ctx.stream()); + ggml_cuda_kernel_launch(comb ? hc_injection_post_f32 : hc_injection_post_f32, launch, a, g); + return; + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const size_t shared = (32 + (a.ncols + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_injection_norm_kernel<256, true, true>(layout) : hc_injection_norm_kernel<256, true, false>(layout)) : (mul ? hc_injection_norm_kernel<256, false, true>(layout) : hc_injection_norm_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_injection_norm_kernel<1024, true, true>(layout) : hc_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_injection_norm_kernel<1024, false, true>(layout) : hc_injection_norm_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, g); +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 95618df5fd4..242911097ff 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -18,3 +18,9 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * mul); + +void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); + +void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 84788c5dc5a..a90f7df786d 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -199,6 +199,108 @@ void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } +struct affine_unary_tensor { + const void * data; + ggml_type type; + uint32_t ne[GGML_MAX_DIMS]; + size_t nb[GGML_MAX_DIMS]; +}; + +struct affine_unary_args { + affine_unary_tensor input[3]; + void * output[4]; + ggml_type output_type[4]; + uint32_t ne[GGML_MAX_DIMS]; + int count; + ggml_unary_op op; + float scale; + float bias; +}; + +static __device__ __forceinline__ float affine_unary_read(const affine_unary_tensor & tensor, const uint32_t * pos) { + size_t offset = 0; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + offset += size_t(pos[d] % tensor.ne[d])*tensor.nb[d]; + } + const char * ptr = (const char *) tensor.data + offset; + if (tensor.type == GGML_TYPE_F16) { + return ggml_cuda_cast(*(const half *) ptr); + } + if (tensor.type == GGML_TYPE_BF16) { + return ggml_cuda_cast(*(const nv_bfloat16 *) ptr); + } + return *(const float *) ptr; +} + +static __device__ __forceinline__ float affine_unary_store(void * dst, ggml_type type, int i, float value) { + if (type == GGML_TYPE_F16) { + const half rounded = ggml_cuda_cast(value); + ((half *) dst)[i] = rounded; + return ggml_cuda_cast(rounded); + } + if (type == GGML_TYPE_BF16) { + const nv_bfloat16 rounded = ggml_cuda_cast(value); + ((nv_bfloat16 *) dst)[i] = rounded; + return ggml_cuda_cast(rounded); + } + ((float *) dst)[i] = value; + return value; +} + +static __global__ void affine_unary_kernel(affine_unary_args args) { + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= args.count) { + return; + } + uint32_t remaining = i; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + pos[d] = remaining % args.ne[d]; + remaining /= args.ne[d]; + } + ggml_cuda_pdl_sync(); + float value = __fmul_rn(affine_unary_read(args.input[0], pos), affine_unary_read(args.input[1], pos)); + value = affine_unary_store(args.output[0], args.output_type[0], i, value); + value = __fadd_rn(value, affine_unary_read(args.input[2], pos)); + value = affine_unary_store(args.output[1], args.output_type[1], i, value); + value = args.op == GGML_UNARY_OP_SIGMOID ? op_sigmoid(value) : op_silu(value); + value = affine_unary_store(args.output[2], args.output_type[2], i, value); + if (args.output[3]) { + affine_unary_store(args.output[3], args.output_type[3], i, args.scale*value + args.bias); + } +} + +void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { + affine_unary_args args{}; + const ggml_tensor * input[] = { mul->src[0], mul->src[1], add->src[1] }; + const ggml_tensor * output[] = { mul, add, unary, post }; + for (int j = 0; j < 3; ++j) { + args.input[j].data = input[j]->data; + args.input[j].type = input[j]->type; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + args.input[j].ne[d] = input[j]->ne[d]; + args.input[j].nb[d] = input[j]->nb[d]; + } + } + for (int j = 0; j < 4; ++j) { + if (output[j]) { + args.output[j] = output[j]->data; + args.output_type[j] = output[j]->type; + } + args.ne[j] = mul->ne[j]; + } + args.count = ggml_nelements(mul); + args.op = ggml_get_unary_op(unary); + if (post) { + memcpy(&args.scale, post->op_params, sizeof(float)); + memcpy(&args.bias, (const char *) post->op_params + sizeof(float), sizeof(float)); + } + const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(affine_unary_kernel, launch, args); +} + void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 04f3af6443a..8439d85afeb 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -93,6 +93,8 @@ void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); +void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post); + void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); __device__ __forceinline__ float ggml_cuda_op_silu_single(float x) { From aed72ffeae08651e420770cef1e3311bbefc2fa9 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 01:55:29 -0400 Subject: [PATCH 02/27] cuda : fuse tensor affine gates with HC postprocessing Preserve gate intermediates, original rounding and live HC outputs. Retain strict input alias checks and standard allocation dependencies. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 166 ++++++++++++++++++++++++ ggml/src/ggml-cuda/norm.cu | 218 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/norm.cuh | 2 + 3 files changed, 386 insertions(+) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a7687b61c92..fe28a8c8e36 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4019,6 +4019,143 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu return match; } +struct ggml_cuda_hc_affine_injection_match { + ggml_cuda_affine_unary_match affine; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + ggml_cuda_hc_post_norm_match hc; + int count = 0; +}; + +static bool ggml_cuda_hc_affine_tensor_ok(const ggml_tensor * tensor) { + if (!tensor || tensor->type != GGML_TYPE_F32 || (tensor->view_src && tensor->view_src->type != GGML_TYPE_F32)) { return false; } + size_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0 || tensor->ne[d] > INT_MAX || tensor->nb[d] % sizeof(float) || elements > INT64_MAX/sizeof(float)/size_t(tensor->ne[d])) { return false; } + elements *= tensor->ne[d]; + } + size_t bytes; + return ggml_cuda_affine_unary_bytes(tensor, bytes) && bytes <= INT64_MAX; +} + +static bool ggml_cuda_hc_affine_prior_source(ggml_cgraph * graph, int start, const ggml_tensor * tensor) { + if (!tensor) { return false; } + if (tensor->view_src) { return ggml_cuda_affine_unary_prior_view(graph, start, tensor); } + if (tensor->op == GGML_OP_NONE) { return true; } + for (int i = 0; i < start; ++i) { + if (graph->nodes[i] == tensor) { return true; } + } + return false; +} + +static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + auto affine = ggml_cuda_match_affine_unary(ctx, graph, i, allocated); + if (!affine.count || affine.mul->type != GGML_TYPE_F32 || affine.add->type != GGML_TYPE_F32 || affine.unary->type != GGML_TYPE_F32) { return {}; } + ggml_tensor * last = affine.post ? affine.post : affine.unary; + const ggml_tensor * gate_inputs[] = {affine.mul->src[0], affine.mul->src[1], affine.add->src[1]}; + for (const ggml_tensor * tensor : gate_inputs) { + if (!ggml_cuda_hc_affine_tensor_ok(tensor) || tensor->ne[2] != 1 || tensor->ne[3] != 1 || affine.mul->ne[0] % tensor->ne[0] || affine.mul->ne[1] % tensor->ne[1]) { return {}; } + } + int post_i = i + affine.count; + auto skip_views = [&]() { + while (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_VIEW) { + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[post_i])) { return false; } + ++post_i; + } + return true; + }; + if (!skip_views()) { return {}; } + ggml_tensor * comb = nullptr; + if (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_DSV4_HC_COMB) { + comb = graph->nodes[post_i++]; + // The existing COMB operator supports four streams. + if (comb->type != GGML_TYPE_F32 || comb->ne[0] != 4 || comb->ne[1] != 4 || comb->ne[3] != 1) { return {}; } + const int64_t mix_dim = (comb->ne[0] + 2)*comb->ne[0]; + const ggml_tensor * mixes = comb->src[0]; + const ggml_tensor * scale = comb->src[1]; + const ggml_tensor * base = comb->src[2]; + if (!mixes || !scale || !base || mixes->ne[0] != mix_dim || mixes->ne[1] != comb->ne[2] || mixes->ne[2] != 1 || mixes->ne[3] != 1 || + scale->ne[0] < 3 || scale->ne[1] != 1 || scale->ne[2] != 1 || scale->ne[3] != 1 || + base->ne[0] != mix_dim || base->ne[1] != 1 || base->ne[2] != 1 || base->ne[3] != 1 || comb->op_params[1] <= 0) { return {}; } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (comb->src[j] && (!ggml_cuda_hc_affine_tensor_ok(comb->src[j]) || !ggml_cuda_hc_affine_prior_source(graph, i, comb->src[j]))) { return {}; } + } + if (!skip_views()) { return {}; } + } + if (post_i >= graph->n_nodes) { return {}; } + ggml_tensor * post = graph->nodes[post_i]; + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * post_comb = post->src[3]; + if (post->op != GGML_OP_DSV4_HC_POST || post->src[2] != last || !x || !residual || + x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || + !ggml_is_contiguous(post) || !ggml_are_same_shape(post, residual) || x->ne[0] <= 0 || x->ne[1] <= 0 || residual->ne[1] <= 0 || + x->ne[2] != 1 || x->ne[3] != 1 || residual->ne[0] != x->ne[0] || residual->ne[2] != x->ne[1] || residual->ne[3] != 1 || + last->ne[0] != residual->ne[1] || last->ne[1] != x->ne[1] || last->ne[2] != 1 || last->ne[3] != 1 || + (comb && post_comb != comb) || + (post_comb && (post_comb->type != GGML_TYPE_F32 || post_comb->ne[0] != residual->ne[1] || post_comb->ne[1] != residual->ne[1] || post_comb->ne[2] != x->ne[1] || post_comb->ne[3] != 1))) { return {}; } + for (const ggml_tensor * tensor : {x, residual, (const ggml_tensor *) post, post_comb}) { + if (tensor && !ggml_cuda_hc_affine_tensor_ok(tensor)) { return {}; } + } + size_t bytes; + if (!ggml_cuda_hc_injection_bytes(post, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } + auto hc = ggml_cuda_match_hc_post_norm(graph, post_i); + if (hc.count && hc.norm->ne[0] != x->ne[0]) { + if (post_comb && bytes/sizeof(float) <= 256) { return {}; } + if (hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] > WARP_SIZE/2 && hc.norm->ne[0] < WARP_SIZE) { return {}; } + const int block_size = hc.norm->ne[0] < 1024 ? 256 : 1024; + const int64_t nweights = (hc.norm->ne[0] + x->ne[0] - 2)/x->ne[0] + 1; + if ((hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] <= WARP_SIZE/2) || nweights >= block_size) { hc = {}; } + } + const int end = hc.count ? post_i + hc.count : post_i + 1; + const int64_t ncols = hc.count ? hc.norm->ne[0] : 256; + const size_t shared = ((hc.count ? 32 : 0) + (ncols + x->ne[0] - 2)/x->ne[0] + 1)*sizeof(float); + if (shared > ggml_cuda_info().devices[ctx.device].smpb) { return {}; } + std::vector ops; + std::vector outputs; + for (int j = i; j < end; ++j) { + ops.push_back(graph->nodes[j]->op); + if (!hc.mul || graph->nodes[j] != hc.norm) { outputs.push_back(j); } + } + if (!ggml_can_fuse_subgraph(graph, i, end - i, ops.data(), outputs.data(), outputs.size())) { return {}; } + if (allocated) { + if (hc.count && !ggml_cuda_hc_post_norm_memory_ok(graph, post_i, hc, ctx.device)) { return {}; } + std::vector reads = {gate_inputs[0], gate_inputs[1], gate_inputs[2], x, residual, post_comb}; + if (hc.mul) { reads.push_back(hc.mul->src[hc.mul->src[0] == hc.norm ? 1 : 0]); } + if (comb) { + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (comb->src[j]) { reads.push_back(comb->src[j]); } + } + } + std::vector writes = {affine.mul, affine.add, affine.unary}; + if (affine.post) { writes.push_back(affine.post); } + writes.push_back(post); + if (hc.count) { writes.push_back(hc.mul ? hc.mul : hc.norm); } + if (comb) { writes.push_back(comb); } + std::vector rb(reads.size()), re(reads.size()), wb(writes.size()), we(writes.size()); + for (size_t j = 0; j < reads.size(); ++j) { + if (reads[j] && (!ggml_cuda_hc_affine_tensor_ok(reads[j]) || !ggml_cuda_hc_injection_range(reads[j], ctx.device, rb[j], re[j]))) { return {}; } + } + for (size_t j = 0; j < writes.size(); ++j) { + if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_is_contiguous(writes[j]) || !ggml_cuda_hc_affine_tensor_ok(writes[j]) || !ggml_cuda_hc_injection_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } + for (size_t k = 0; k < reads.size(); ++k) { + // COMB is written before the affine kernel reads it. + if (reads[k] && !(writes[j] == comb && reads[k] == comb) && wb[j] < re[k] && rb[k] < we[j]) { return {}; } + } + for (size_t k = 0; k < j; ++k) { + if (wb[j] < we[k] && wb[k] < we[j]) { return {}; } + } + } + for (int j = i; j < end; ++j) { + const ggml_tensor * tensor = graph->nodes[j]; + if (tensor->op != GGML_OP_VIEW && tensor->op != GGML_OP_RESHAPE) { continue; } + uintptr_t begin, stop; + if (!ggml_cuda_hc_injection_range(tensor, ctx.device, begin, stop)) { return {}; } + } + } + return {affine, post, comb, hc, end - i}; +} + // 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) { @@ -4792,6 +4929,19 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); + if (injection.count) { + if (injection.comb) { + const bool ok = ggml_cuda_compute_forward(*cuda_ctx, injection.comb); + GGML_ASSERT(ok); + } + const auto & affine = injection.affine; + ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, injection.post, injection.hc.norm, injection.hc.mul); + return injection.count - 1; + } + } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); if (affine.count) { @@ -5152,6 +5302,22 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { + const auto affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i, false); + if (affine_injection.count) { + ggml_tensor * terminal = cgraph->nodes[i + affine_injection.count - 1]; + ggml_tensor * unused_norm = affine_injection.hc.mul ? affine_injection.hc.norm : nullptr; + for (int j = i; j < i + affine_injection.count - 1; ++j) { + if (cgraph->nodes[j] != unused_norm) { params->add_alloc_dep(params->user_data, cgraph->nodes[j], terminal); } + for (int k = 0; k < GGML_MAX_SRC; ++k) { + if (cgraph->nodes[j]->src[k]) { params->add_alloc_dep(params->user_data, cgraph->nodes[j]->src[k], terminal); } + } + } + for (ggml_tensor * src : terminal->src) { + if (src && src != unused_norm) { params->add_alloc_dep(params->user_data, src, terminal); } + } + i += affine_injection.count - 1; + continue; + } const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); if (affine.count) { ggml_tensor * last = cgraph->nodes[i + affine.count - 1]; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index e5b6f2b18f8..27d3c4c0f67 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -912,6 +912,49 @@ static __device__ __forceinline__ float hc_injection_gate(hc_injection_data g, i return value; } +struct hc_affine_injection_operand { + const float * data; + uint3 cols; + uint3 rows; + int64_t s0; + int64_t s1; +}; + +struct hc_affine_injection_data { + hc_affine_injection_operand input; + hc_affine_injection_operand multiplier; + hc_affine_injection_operand bias; + uint3 streams; + float * product; + float * added; + float * unary; + float * last; + float scale1; + float bias1; + ggml_unary_op op; + int64_t values; +}; + +static __device__ __forceinline__ float hc_affine_injection_read(hc_affine_injection_operand a, uint32_t col, uint32_t row) { + return a.data[int64_t(fastmodulo(col, a.cols))*a.s0 + int64_t(fastmodulo(row, a.rows))*a.s1]; +} + +static __device__ __forceinline__ float hc_injection_gate(hc_affine_injection_data g, int64_t index, bool store) { + const uint32_t col = fastmodulo(uint32_t(index), g.streams); + const uint32_t row = fastdiv(uint32_t(index), g.streams); + float value = __fmul_rn(hc_affine_injection_read(g.input, col, row), hc_affine_injection_read(g.multiplier, col, row)); + if (store) { g.product[index] = value; } + value = __fadd_rn(value, hc_affine_injection_read(g.bias, col, row)); + if (store) { g.added[index] = value; } + value = g.op == GGML_UNARY_OP_SILU ? ggml_cuda_op_silu_single(value) : 1.0f / (1.0f + expf(-value)); + if (store) { g.unary[index] = value; } + if (g.last) { + value = __fmaf_rn(g.scale1, value, g.bias1); + if (store) { g.last[index] = value; } + } + return value; +} + static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { const int64_t scalar = first + j; @@ -920,6 +963,99 @@ static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, fl __syncthreads(); } +template +static __device__ __forceinline__ void hc_injection_load(Gate g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { + for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { + const int64_t scalar = first + j; + weights[j] = hc_injection_gate(g, scalar, scalar*n_embd >= begin); + } + __syncthreads(); +} + +template +static __device__ __forceinline__ void hc_injection_norm_impl(hc_post_norm_data a, Gate g) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + ggml_cuda_pdl_sync(); + extern __shared__ float s_sum[]; + float * weights = s_sum + 32; + const int64_t gate_first = offset/a.n_embd; + float grouped_weight = 0.0f; + if constexpr (grouped) { + if (threadIdx.x % WARP_SIZE == 0) { + grouped_weight = hc_injection_gate(g, gate_first, threadIdx.x == 0); + } + grouped_weight = __shfl_sync(0xffffffff, grouped_weight, 0, WARP_SIZE); + } else { + const int64_t gate_count = (offset + a.ncols - 1)/a.n_embd - gate_first + 1; + hc_injection_load(g, weights, gate_first, gate_count, offset, a.n_embd); + } + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; + } + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + const float weight = grouped ? grouped_weight : weights[ih + it*a.hc - gate_first]; + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weight; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + a.dst[offset + col] = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + } else { + a.dst[offset + col] = scale * a.post_dst[offset + col]; + } + } +} + +template +static __device__ __forceinline__ void hc_injection_post_impl(hc_post_norm_data a, Gate g) { + ggml_cuda_pdl_lc(); + const int64_t begin = int64_t(blockIdx.x)*blockDim.x; + const int64_t end = begin + int64_t(blockDim.x) < g.values ? begin + int64_t(blockDim.x) : g.values; + const int64_t ir = begin + threadIdx.x; + ggml_cuda_pdl_sync(); + extern __shared__ float weights[]; + const int64_t first = begin/a.n_embd; + const int64_t count = (end - 1)/a.n_embd - first + 1; + hc_injection_load(g, weights, first, count, begin, a.n_embd); + if (ir >= g.values) { return; } + const int64_t i0 = ir % a.n_embd; + const int64_t ih = (ir/a.n_embd) % a.hc; + const int64_t it = ir/(a.n_embd*a.hc); + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weights[ir/a.n_embd - first]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; +} + template static __global__ __launch_bounds__(block_size) void hc_injection_norm_f32(hc_post_norm_data a, hc_injection_data g) { ggml_cuda_pdl_lc(); @@ -1004,6 +1140,16 @@ static __global__ void hc_injection_post_f32(hc_post_norm_data a, hc_injection_d a.post_dst[ir] = value; } +template +static __global__ __launch_bounds__(block_size) void hc_affine_injection_norm_f32(hc_post_norm_data a, hc_affine_injection_data g) { + hc_injection_norm_impl(a, g); +} + +template +static __global__ void hc_affine_injection_post_f32(hc_post_norm_data a, hc_affine_injection_data g) { + hc_injection_post_impl(a, g); +} + template static auto hc_injection_norm_kernel(int layout) { if (layout == 1) { return hc_injection_norm_f32; } @@ -1069,3 +1215,75 @@ void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * fi : (comb ? (mul ? hc_injection_norm_kernel<1024, true, true>(layout) : hc_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_injection_norm_kernel<1024, false, true>(layout) : hc_injection_norm_kernel<1024, false, false>(layout))); ggml_cuda_kernel_launch(kernel, launch, a, g); } + +static hc_affine_injection_operand ggml_cuda_hc_affine_injection_operand(const ggml_tensor * tensor) { + return {(const float *) tensor->data, init_fastdiv_values(tensor->ne[0]), init_fastdiv_values(tensor->ne[1]), int64_t(tensor->nb[0]/sizeof(float)), int64_t(tensor->nb[1]/sizeof(float))}; +} + +template +static auto hc_affine_injection_norm_kernel(int layout) { + if (layout == 1) { return hc_affine_injection_norm_f32; } + if (layout == 2) { return hc_affine_injection_norm_f32; } + return hc_affine_injection_norm_f32; +} + +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = norm ? (float *) (mul ? mul : norm)->data : nullptr; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm ? norm->ne[0] : 256; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = norm ? ggml_get_op_params_f32(norm, 0) : 0.0f; + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + hc_affine_injection_data g{}; + g.input = ggml_cuda_hc_affine_injection_operand(first->src[0]); + g.multiplier = ggml_cuda_hc_affine_injection_operand(first->src[1]); + g.bias = ggml_cuda_hc_affine_injection_operand(added->src[1]); + g.streams = init_fastdiv_values(residual->ne[1]); + g.product = (float *) first->data; + g.added = (float *) added->data; + g.unary = (float *) unary->data; + g.last = last != unary ? (float *) last->data : nullptr; + g.scale1 = g.last ? ggml_get_op_params_f32(last, 0) : 1.0f; + g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; + g.op = ggml_get_unary_op(unary); + g.values = ggml_nelements(post); + if (!norm) { + const size_t shared = ((256 + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const ggml_cuda_kernel_launch_params launch(dim3((g.values + 255)/256), dim3(256), shared, ctx.stream()); + ggml_cuda_kernel_launch(comb ? hc_affine_injection_post_f32 : hc_affine_injection_post_f32, launch, a, g); + return; + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const size_t shared = (32 + (a.ncols + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_affine_injection_norm_kernel<256, true, true>(layout) : hc_affine_injection_norm_kernel<256, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<256, false, true>(layout) : hc_affine_injection_norm_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_affine_injection_norm_kernel<1024, true, true>(layout) : hc_affine_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<1024, false, true>(layout) : hc_affine_injection_norm_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, g); +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 242911097ff..a2bf28bfa16 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -24,3 +24,5 @@ bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_ten void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); + +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); From 442e7241f76bd60c8f6210edb7014b2c9ff0490c Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 02:31:16 -0400 Subject: [PATCH 03/27] cuda : compose HC prepared-image review prerequisites Use the exact tested control sources from published PR130, PR129 and PR132. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 1512 ++++++++++++++++++++++++++----- ggml/src/ggml-cuda/mmq.cu | 105 ++- ggml/src/ggml-cuda/mmq.cuh | 17 +- ggml/src/ggml-cuda/mmvq.cu | 24 +- ggml/src/ggml-cuda/mmvq.cuh | 5 +- ggml/src/ggml-cuda/norm.cu | 1286 +++++++++++++++++++++----- ggml/src/ggml-cuda/norm.cuh | 25 +- ggml/src/ggml-cuda/reuse.cuh | 194 ++++ ggml/src/ggml-cuda/unary.cu | 72 ++ ggml/src/ggml-cuda/unary.cuh | 2 + 10 files changed, 2700 insertions(+), 542 deletions(-) create mode 100644 ggml/src/ggml-cuda/reuse.cuh diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index fe28a8c8e36..625f1595af3 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4,6 +4,7 @@ #include "ggml-cuda/allreduce.cuh" #include "ggml-cuda/common.cuh" +#include "ggml-cuda/reuse.cuh" #include "ggml-cuda/acc.cuh" #include "ggml-cuda/add-id.cuh" #include "ggml-cuda/arange.cuh" @@ -1401,7 +1402,7 @@ struct batched_mul_mat_traits { }; template -static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr) { using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; @@ -1464,7 +1465,16 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const src0_ptr = src0_alloc.get(); } - if (src1->type == compute_type) { + if (prepared_src1) { + GGML_ASSERT(src1->type != compute_type); + src1_ptr = (const cuda_t *) prepared_src1; + if (!ggml_is_contiguously_allocated(src1)) { + s11 = ne10; + s12 = ne11*s11; + s13 = ne12*s12; + is_src1_cont_2 = true; + } + } else if (src1->type == compute_type) { src1_ptr = (const cuda_t *) src1->data; } else { src1_alloc.alloc(ggml_nelements(src1)); @@ -1614,8 +1624,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const } } -static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - const int cc = ggml_cuda_info().devices[ctx.device].cc; +static ggml_type ggml_cuda_mul_mat_cublas_compute_type(int cc, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, bool warn = true) { const ggml_prec prec = (ggml_prec) ggml_get_op_params_i32(dst, 0); ggml_type compute_type = src0->type; if (ggml_is_quantized(compute_type)) { @@ -1649,20 +1658,25 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml compute_type = GGML_TYPE_F16; } else if (env_cpp == "bf16") { compute_type = GGML_TYPE_BF16; - } else if (env_cpp != "auto") { - GGML_LOG_WARN("%s: unknown value for GGML_CUDA_CUBLAS_COMPUTE_TYPE: %s", __func__, env_cpp.c_str()); + } else if (warn && env_cpp != "auto") { + GGML_LOG_WARN("%s: unknown value for GGML_CUDA_CUBLAS_COMPUTE_TYPE: %s", "ggml_cuda_mul_mat_cublas", env_cpp.c_str()); } } + return compute_type; +} + +static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr) { + const ggml_type compute_type = ggml_cuda_mul_mat_cublas_compute_type(ggml_cuda_info().devices[ctx.device].cc, src0, src1, dst, true); switch (compute_type) { case GGML_TYPE_F32: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); break; case GGML_TYPE_BF16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); break; case GGML_TYPE_F16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); break; default: GGML_ABORT("fatal error"); @@ -1819,59 +1833,84 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -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 - - if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) { - return; - } +enum ggml_cuda_mm_kernel { + GGML_CUDA_MM_MMVF, + GGML_CUDA_MM_MMVF_TRANSPOSE, + GGML_CUDA_MM_MMF, + GGML_CUDA_MM_MMVQ, + GGML_CUDA_MM_MMQ, + GGML_CUDA_MM_CUBLAS, +}; +static ggml_cuda_mm_kernel ggml_cuda_mul_mat_kernel(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int device) { + 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. // Therefore, in such cases use cuBLAS. 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_MM_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_MM_MMVF; } // 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_MM_MMVF_TRANSPOSE; } 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_MM_MMF; } if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVQ; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return GGML_CUDA_MM_MMQ; + } + return GGML_CUDA_MM_CUBLAS; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr) { + if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + switch (ggml_cuda_mul_mat_kernel(src0, src1, dst, ctx.device)) { + case GGML_CUDA_MM_MMVF: + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + break; + case GGML_CUDA_MM_MMVF_TRANSPOSE: { + ggml_tensor dst_vec = *dst; + dst_vec.ne[0] = src1->ne[1]; + dst_vec.ne[1] = 1; + dst_vec.nb[1] = dst_vec.nb[0]*src1->ne[1]; + 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); + break; + } + case GGML_CUDA_MM_MMF: + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + break; + case GGML_CUDA_MM_MMVQ: + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + break; + case GGML_CUDA_MM_MMQ: + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + break; + case GGML_CUDA_MM_CUBLAS: + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst, prepared_src1); + break; + } } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization @@ -2063,7 +2102,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * nb1, nb2, nb3, stream); } -static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { +static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst, const void * prepared_src1 = nullptr) { switch (dst->op) { case GGML_OP_ARGMAX: ggml_cuda_argmax(ctx, dst); @@ -2256,7 +2295,7 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg ggml_cuda_op_rms_norm_back(ctx, dst); break; case GGML_OP_MUL_MAT: - ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); + ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst, prepared_src1); break; case GGML_OP_MUL_MAT_ID: ggml_cuda_mul_mat_id(ctx, dst); @@ -3034,6 +3073,58 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, return is_ok; } +static bool ggml_cuda_can_share_mmq_input(const ggml_tensor * node, int device) { + if (node->op != GGML_OP_MUL_MAT || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || node->view_src || + node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node)) { + return false; + } + const ggml_tensor * weight = node->src[0]; + const ggml_tensor * input = node->src[1]; + const int cc = ggml_cuda_info().devices[device].cc; + if (input->type != GGML_TYPE_F32 || input->nb[0] != sizeof(float) || input->ne[1] < 1 || + input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || + !ggml_is_contiguous(weight) || ggml_cuda_should_use_mmvq(weight->type, cc, input->ne[1]) || + !ggml_cuda_should_use_mmq(weight->type, cc, input->ne[1], /*n_experts =*/ 0) || ggml_cuda_op_mul_mat_use_fwht(node)) { + return false; + } + for (const ggml_tensor * tensor : { weight, input, node }) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + } + return !(weight->view_src && ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(weight) != ggml_backend_buffer_get_alloc_size(weight->buffer, weight)); +} + +static bool ggml_cuda_mmq_input_overwritten(const ggml_tensor * node, const ggml_tensor * input) { + if (node->op == GGML_OP_OPT_STEP_ADAMW || node->op == GGML_OP_OPT_STEP_SGD) { + return true; + } + if (ggml_cuda_is_view_or_noop(node) || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return false; + } + const uintptr_t start = (uintptr_t) input->data; + const uintptr_t end = start + ggml_backend_buffer_get_alloc_size(input->buffer, input); + const auto overlaps = [&](const void * data, size_t size) { + const uintptr_t other = (uintptr_t) data; + return start < other + size && other < end; + }; + if (!node->buffer || !node->data || overlaps(node->data, ggml_backend_buffer_get_alloc_size(node->buffer, node))) { + return true; + } + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const ggml_tensor * weight = node->src[0]; + if (weight->buffer && ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size = ggml_nbytes(weight); + const size_t alloc = ggml_backend_buffer_get_alloc_size(weight->buffer, weight); + if (alloc > size && overlaps((const char *) weight->data + size, alloc - size)) { + return true; + } + } + } + return false; +} + // The long form spans 2*k + 1 nodes. ggml_can_fuse_subgraph() accepts at most // 31 nodes, so k <= 15; larger values use the per-operation path. static constexpr int MOE_WEIGHTED_REDUCTION_MAX_EXPERTS = 15; @@ -3445,10 +3536,238 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } +static bool ggml_cuda_can_share_mmvq_input(const ggml_tensor * node, int device) { + if (node->op != GGML_OP_MUL_MAT || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || node->view_src || + node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node)) { + return false; + } + const ggml_tensor * weight = node->src[0]; + const ggml_tensor * input = node->src[1]; + const int cc = ggml_cuda_info().devices[device].cc; + if (cc <= GGML_CUDA_CC_PASCAL || input->type != GGML_TYPE_F32 || input->nb[0] != sizeof(float) || + input->ne[1] < 1 || input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || + !ggml_is_contiguous(weight) || !ggml_cuda_should_use_mmvq(weight->type, cc, input->ne[1]) || + ggml_cuda_op_mul_mat_use_fwht(node)) { + return false; + } + for (const ggml_tensor * tensor : { weight, input, node }) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + } + return !(weight->view_src && ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(weight) != ggml_backend_buffer_get_alloc_size(weight->buffer, weight)); +} + +static bool ggml_cuda_mmvq_input_overwritten(const ggml_tensor * node, const ggml_tensor * input) { + if (node->op == GGML_OP_OPT_STEP_ADAMW || node->op == GGML_OP_OPT_STEP_SGD) { + return true; + } + if (ggml_cuda_is_view_or_noop(node) || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return false; + } + const uintptr_t start = (uintptr_t) input->data; + const uintptr_t end = start + ggml_backend_buffer_get_alloc_size(input->buffer, input); + const auto overlaps = [&](const void * data, size_t size) { + const uintptr_t other = (uintptr_t) data; + return start < other + size && other < end; + }; + if (!node->buffer || !node->data || overlaps(node->data, ggml_backend_buffer_get_alloc_size(node->buffer, node))) { + return true; + } + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const ggml_tensor * weight = node->src[0]; + if (weight->buffer && ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size = ggml_nbytes(weight); + const size_t alloc = ggml_backend_buffer_get_alloc_size(weight->buffer, weight); + if (alloc > size && overlaps((const char *) weight->data + size, alloc - size)) { + return true; + } + } + } + return false; +} + +static bool ggml_cuda_prepared_bytes(const ggml_tensor * tensor, size_t & bytes) { + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0) { return false; } + } + const size_t block = ggml_blck_size(tensor->type); + bytes = ggml_type_size(tensor->type); + int first = 0; + if (block > 1) { + if (!tensor->nb[0] || uint64_t(tensor->ne[0]) > SIZE_MAX/tensor->nb[0] || tensor->ne[0] % block) { + return false; + } + bytes = size_t(tensor->ne[0])*tensor->nb[0]/block; + first = 1; + } + for (int d = first; d < GGML_MAX_DIMS; ++d) { + const uint64_t count = tensor->ne[d] - 1; + if (count && tensor->nb[d] > (SIZE_MAX - bytes)/count) { return false; } + bytes += size_t(count)*tensor->nb[d]; + } + return bytes && bytes == ggml_nbytes(tensor); +} + +static bool ggml_cuda_prepared_range(const ggml_tensor * tensor, int device, uintptr_t & begin, uintptr_t & end) { + if (!tensor->buffer || !tensor->data || tensor->buffer->buft != ggml_backend_cuda_buffer_type(device)) { + return false; + } + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + size_t bytes; + if (!ggml_cuda_prepared_bytes(tensor, bytes)) { return false; } + begin = (uintptr_t) tensor->data; + if (begin < base || begin - base > size || bytes > size - (begin - base) || begin > UINTPTR_MAX - bytes) { + return false; + } + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + const uintptr_t root_begin = (uintptr_t) root->data; + size_t root_bytes; + if (!ggml_cuda_prepared_bytes(root, root_bytes)) { return false; } + if (root->view_src || root->buffer != tensor->buffer || !root->data || root_begin < base || root_begin - base > size || + root_bytes > size - (root_begin - base) || tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs || + root_begin > UINTPTR_MAX - tensor->view_offs || begin != root_begin + tensor->view_offs) { + return false; + } + } + end = begin + bytes; + return true; +} + +static int ggml_cuda_cublas_input_key(const ggml_tensor * node, int device, size_t & size) { + if (node->op != GGML_OP_MUL_MAT || !(node->flags & GGML_TENSOR_FLAG_COMPUTE) || node->view_src || node->type != GGML_TYPE_F32 || + !node->src[0] || !node->src[1] || !ggml_is_contiguous(node) || ggml_cuda_op_mul_mat_use_fwht(node)) { + return -1; + } + const ggml_tensor * weight = node->src[0], * input = node->src[1]; + for (const ggml_tensor * tensor : {weight, input, node}) { + if (tensor != weight && tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16 && tensor->type != GGML_TYPE_BF16) { + return -1; + } + uintptr_t begin, end; + if (!ggml_cuda_prepared_range(tensor, device, begin, end) || tensor->nb[0] != ggml_type_size(tensor->type) || (tensor != weight && begin % ggml_type_size(tensor->type))) { + return -1; + } + } + if (ggml_cuda_mul_mat_kernel(weight, input, node, device) != GGML_CUDA_MM_CUBLAS) { + return -1; + } + const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(ggml_cuda_info().devices[device].cc, weight, input, node, false); + if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || type == input->type) { + return -1; + } + size_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (input->ne[d] <= 0 || uint64_t(input->ne[d]) > SIZE_MAX/elements/sizeof(half)) { + return -1; + } + elements *= input->ne[d]; + } + size = elements*sizeof(half); + return int(type); +} + +static bool ggml_cuda_prepared_input_overwritten(const ggml_tensor * node, const ggml_tensor * input) { + if (node->op == GGML_OP_OPT_STEP_ADAMW || node->op == GGML_OP_OPT_STEP_SGD) { + return true; + } + if (ggml_cuda_is_view_or_noop(node) || !(node->flags & GGML_TENSOR_FLAG_COMPUTE)) { + return false; + } + const int device = ((const ggml_backend_cuda_buffer_context *) input->buffer->context)->device; + uintptr_t begin, end, other_begin, other_end; + if (!ggml_cuda_prepared_range(input, device, begin, end) || !ggml_cuda_prepared_range(node, device, other_begin, other_end)) { + return true; + } + if (node->op == GGML_OP_FLASH_ATTN_EXT) { + size_t maximum = other_end - other_begin; + for (int j = 1; j <= 2; ++j) { + const ggml_tensor * tensor = node->src[j]; + size_t bytes, elements = 1; + if (!tensor || !ggml_cuda_prepared_bytes(tensor, bytes)) { return true; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (uint64_t(tensor->ne[d]) > INT64_MAX/elements) { return true; } + elements *= tensor->ne[d]; + } + if (maximum > SIZE_MAX - 127 || elements > (SIZE_MAX - 127 - maximum)/sizeof(half)) { return true; } + maximum += 127 + elements*sizeof(half); + } + if (other_begin > UINTPTR_MAX - maximum) { return true; } + } + const size_t allocated = ggml_backend_buffer_get_alloc_size(node->buffer, node); + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(node->buffer); + const size_t size = ggml_backend_buffer_get_size(node->buffer); + if (allocated < other_end - other_begin || allocated > size - (other_begin - base) || other_begin > UINTPTR_MAX - allocated || + (begin < other_begin + allocated && other_begin < end)) { + return true; + } + if (node->op == GGML_OP_SET) { + if (!node->src[1]) { return true; } + ggml_tensor slice = *node; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { slice.ne[d] = node->src[1]->ne[d]; } + slice.nb[0] = ggml_type_size(node->type); + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + if (node->op_params[d - 1] < 0) { return true; } + slice.nb[d] = node->op_params[d - 1]; + } + size_t bytes; + if (node->op_params[3] < 0 || !ggml_cuda_prepared_bytes(&slice, bytes) || + size_t(node->op_params[3]) > allocated || bytes > allocated - size_t(node->op_params[3])) { + return true; + } + } + if (node->op == GGML_OP_CPY) { + if (!node->src[1] || !ggml_cuda_prepared_range(node->src[1], device, other_begin, other_end) || + (begin < other_end && other_begin < end)) { + return true; + } + } + if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { + const ggml_tensor * weight = node->src[0]; + if (weight->buffer && ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + size_t bytes; + if (!ggml_cuda_prepared_bytes(weight, bytes)) { return true; } + const size_t allocated = ggml_backend_buffer_get_alloc_size(weight->buffer, weight); + if (allocated > bytes) { + if (!ggml_cuda_prepared_range(weight, device, other_begin, other_end) || other_begin > UINTPTR_MAX - allocated || + (begin < other_begin + allocated && other_end < end)) { + return true; + } + } + } + } + return false; +} + +static void ggml_cuda_prepare_cublas_input(ggml_backend_cuda_context & ctx, const ggml_tensor * input, int key, ggml_cuda_pool_alloc & prepared) { + const size_t elements = ggml_nelements(input); + prepared.alloc(elements*sizeof(half)); + const bool contiguous = ggml_is_contiguously_allocated(input); + if (key == GGML_TYPE_BF16) { + if (contiguous) { + ggml_get_to_bf16_cuda(input->type)(input->data, (nv_bfloat16 *) prepared.get(), elements, ctx.stream()); + } else { + ggml_get_to_bf16_nc_cuda(input->type)(input->data, (nv_bfloat16 *) prepared.get(), input->ne[0], input->ne[1], input->ne[2], input->ne[3], + input->nb[1]/ggml_type_size(input->type), input->nb[2]/ggml_type_size(input->type), input->nb[3]/ggml_type_size(input->type), ctx.stream()); + } + } else { + if (contiguous) { + ggml_get_to_fp16_cuda(input->type)(input->data, (half *) prepared.get(), elements, ctx.stream()); + } else { + ggml_get_to_fp16_nc_cuda(input->type)(input->data, (half *) prepared.get(), input->ne[0], input->ne[1], input->ne[2], input->ne[3], + input->nb[1]/ggml_type_size(input->type), input->nb[2]/ggml_type_size(input->type), input->nb[3]/ggml_type_size(input->type), ctx.stream()); + } + } +} + struct ggml_cuda_hc_post_norm_match { ggml_tensor * post = nullptr; ggml_tensor * norm = nullptr; ggml_tensor * mul = nullptr; + ggml_tensor * scale = nullptr; int count = 0; }; @@ -3472,20 +3791,28 @@ static ggml_cuda_hc_post_norm_match ggml_cuda_match_hc_post_norm(ggml_cgraph * c } if (ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {}) || ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {}) || - ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {}) || - ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_SCALE }, {})) { + ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { return match; } + const bool with_scale = ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_SCALE }, {}); const bool weighted = ggml_cuda_can_fuse(cgraph, norm_i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {}); ggml_tensor * mul = weighted ? cgraph->nodes[norm_i + 1] : nullptr; - const bool closed = reshape + ggml_tensor * scale = with_scale ? cgraph->nodes[norm_i + 1] : nullptr; + bool closed = reshape ? (weighted ? ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RESHAPE, GGML_OP_RMS_NORM, GGML_OP_MUL }, { i, i + 1, i + 3 }) : ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RESHAPE, GGML_OP_RMS_NORM }, { i, i + 1, i + 2 })) : (weighted ? ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RMS_NORM, GGML_OP_MUL }, { i, i + 2 }) : ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RMS_NORM }, { i, i + 1 })); - if (closed && ggml_cuda_should_fuse_hc_post_norm(post, norm, mul)) { + if (scale) { + closed = reshape + ? ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RESHAPE, GGML_OP_RMS_NORM, GGML_OP_SCALE }, { i, i + 1, i + 3 }) + : ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_DSV4_HC_POST, GGML_OP_RMS_NORM, GGML_OP_SCALE }, { i, i + 2 }); + } + const bool supported = scale ? ggml_cuda_should_fuse_hc_post_norm_scale(post, norm, scale) : ggml_cuda_should_fuse_hc_post_norm(post, norm, mul); + if (closed && supported) { match.post = post; match.norm = norm; match.mul = mul; - match.count = norm_i - i + 1 + int(weighted); + match.scale = scale; + match.count = norm_i - i + 1 + int(weighted) + int(with_scale); } return match; } @@ -3517,7 +3844,7 @@ static bool ggml_cuda_hc_post_norm_memory_ok(const ggml_cgraph * cgraph, int i, } return true; }; - const ggml_tensor * writes[] = { match.post, match.mul ? match.mul : match.norm }; + const ggml_tensor * writes[] = { match.post, match.scale ? match.scale : match.mul ? match.mul : match.norm }; uintptr_t begin[2], end[2]; for (int w = 0; w < 2; ++w) { if (!range(writes[w], begin[w], end[w])) { @@ -3569,6 +3896,579 @@ static bool ggml_cuda_hc_post_norm_memory_ok(const ggml_cgraph * cgraph, int i, } // try and fuse nodes and return the number of nodes to skip +struct ggml_cuda_norm_emit_match { + ggml_tensor * post = nullptr; + ggml_tensor * norm = nullptr; + ggml_tensor * mul = nullptr; + ggml_tensor * add = nullptr; + ggml_tensor * scale = nullptr; + ggml_tensor * dst = nullptr; + int last = -1; + int f16 = -1; + int bf16 = -1; +}; + +static ggml_cuda_norm_emit_match ggml_cuda_match_norm_emit(ggml_cgraph * graph, int i, int device) { + if (graph->nodes[i]->op == GGML_OP_DSV4_HC_POST) { + const auto hc = ggml_cuda_match_hc_post_norm(graph, i); + if (!hc.count || !ggml_cuda_hc_post_norm_memory_ok(graph, i, hc, device)) { return {}; } + ggml_cuda_norm_emit_match match; + match.post = hc.post; + match.norm = hc.norm; + match.mul = hc.mul; + match.scale = hc.scale; + match.dst = hc.scale ? hc.scale : hc.mul ? hc.mul : hc.norm; + match.last = i + hc.count - 1; + if (!(match.norm->flags & GGML_TENSOR_FLAG_COMPUTE) || !(match.dst->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + return match; + } + ggml_tensor * norm = graph->nodes[i]; + if (norm->op != GGML_OP_RMS_NORM || !norm->src[0] || norm->type != GGML_TYPE_F32 || norm->src[0]->type != GGML_TYPE_F32 || + !(norm->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS}, {}) || + ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE}, {})) { return {}; } + ggml_cuda_norm_emit_match match; + match.norm = match.dst = norm; + match.last = i; + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { + match.mul = graph->nodes[i + 1]; match.add = match.dst = graph->nodes[i + 2]; match.last = i + 2; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { + match.mul = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_SCALE}, {})) { + match.scale = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } + const ggml_tensor * weight = match.mul ? (match.mul->src[0] == norm ? match.mul->src[1] : match.mul->src[0]) : nullptr; + const ggml_tensor * bias = match.add ? (match.add->src[0] == match.mul ? match.add->src[1] : match.add->src[0]) : nullptr; + uintptr_t begin, end; + if (!(match.dst->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_are_same_shape(norm->src[0], norm) || !ggml_are_same_shape(norm, match.dst) || + !ggml_is_contiguous(match.dst) || !ggml_cuda_prepared_range(match.dst, device, begin, end) || + begin % sizeof(float) || (end - begin)/sizeof(float) > INT_MAX) { return {}; } + const ggml_tensor * reads[] = {norm->src[0], weight, bias}; + for (const ggml_tensor * read : reads) { + if (!read) { continue; } + uintptr_t read_begin, read_end; + if (read->type != GGML_TYPE_F32 || read->nb[0] != sizeof(float) || !ggml_cuda_prepared_range(read, device, read_begin, read_end) || + read_begin % sizeof(float) || (read_end - read_begin)/sizeof(float) > INT_MAX || + (begin < read_end && read_begin < end)) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (read->nb[d] % sizeof(float) || read->nb[d]/sizeof(float) > INT64_MAX || read->ne[d] > INT_MAX) { return {}; } + } + } + if (norm->src[0]->ne[2] > 65535 || norm->src[0]->ne[3] > 65535) { return {}; } + return match; +} + +static std::vector ggml_cuda_plan_norm_emit(ggml_cgraph * graph, int device, + const std::vector & keys, const std::vector & sizes, ggml_cuda_reuse_plan & reuse) { + std::unordered_map> readers; + for (int i = 0; i < graph->n_nodes; ++i) { + if (keys[i] != GGML_TYPE_F16 && keys[i] != GGML_TYPE_BF16) { continue; } + const ggml_tensor * input = graph->nodes[i]->src[1]; + readers[input->view_src ? input->view_src : input].push_back(i); + } + if (readers.empty()) { return {}; } + std::vector emits(graph->n_nodes); + bool moved = false; + std::vector removed(reuse.groups.size(), false); + for (int i = 0; i < graph->n_nodes; ++i) { + auto match = ggml_cuda_match_norm_emit(graph, i, device); + if (!match.norm) { continue; } + const auto can_emit = [&](int node, int prepare) { + const ggml_tensor * input = graph->nodes[node]->src[1]; + const ggml_tensor * root = input->view_src ? input->view_src : input; + if (prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + input->type != GGML_TYPE_F32 || ggml_nbytes(input) != ggml_nbytes(match.dst)) { return false; } + for (int j = match.last + 1; j <= prepare; ++j) { + if (ggml_cuda_prepared_input_overwritten(graph->nodes[j], input)) { return false; } + } + return true; + }; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + auto & group = reuse.groups[g]; + const int key = keys[group.node]; + int & target = key == GGML_TYPE_F16 ? match.f16 : match.bf16; + if (removed[g] || target >= 0 || (key != GGML_TYPE_F16 && key != GGML_TYPE_BF16) || !can_emit(group.node, group.prepare)) { continue; } + target = int(g); + group.prepare = i; + group.after = true; + moved = true; + } + const auto found = readers.find(match.dst); + if (found != readers.end()) { + for (const int node : found->second) { + int & target = keys[node] == GGML_TYPE_F16 ? match.f16 : match.bf16; + if (sizes[node] > SIZE_MAX - 255 || !can_emit(node, node)) { continue; } + const int previous = reuse.nodes.empty() ? -1 : reuse.nodes[node]; + if (previous == target && target >= 0) { continue; } + if (target >= 0) { + auto & group = reuse.groups[target]; + if (sizes[node] != sizes[group.node]) { continue; } + if (previous >= 0) { + const auto & other = reuse.groups[previous]; + if (keys[other.node] != keys[group.node] || !can_emit(other.node, other.prepare)) { continue; } + group.last = std::max(group.last, other.last); + group.size = std::max(group.size, other.size); + for (int & member : reuse.nodes) { + if (member == previous) { member = target; } + } + removed[previous] = true; + } else { + reuse.nodes[node] = target; + group.last = std::max(group.last, node); + group.size = std::max(group.size, GGML_PAD(sizes[node], 256)); + } + } else { + if (previous >= 0) { continue; } + if (reuse.nodes.empty()) { + reuse.nodes.assign(graph->n_nodes, -1); + reuse.starts.assign(graph->n_nodes, -1); + } + target = int(reuse.groups.size()); + reuse.groups.push_back({node, i, node, -1, true, GGML_PAD(sizes[node], 256)}); + removed.push_back(false); + reuse.nodes[node] = target; + } + moved = true; + } + } + if (match.f16 >= 0 || match.bf16 >= 0) { emits[i] = match; } + } + if (moved) { + std::vector remap(reuse.groups.size(), -1); + int next = 0; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + if (removed[g]) { continue; } + remap[g] = next; + reuse.groups[next++] = reuse.groups[g]; + } + reuse.groups.resize(next); + for (int & group : reuse.nodes) { if (group >= 0) { group = remap[group]; } } + for (auto & emit : emits) { + if (emit.f16 >= 0) { emit.f16 = remap[emit.f16]; } + if (emit.bf16 >= 0) { emit.bf16 = remap[emit.bf16]; } + } + reuse.pack(); + } + if (reuse.groups.empty()) { return {}; } + return emits; +} + +struct ggml_cuda_norm_q8_match { + ggml_tensor * post = nullptr; + ggml_tensor * norm = nullptr; + ggml_tensor * mul = nullptr; + ggml_tensor * add = nullptr; + ggml_tensor * scale = nullptr; + ggml_tensor * dst = nullptr; + int last = -1; + int image = -1; +}; + +static ggml_cuda_norm_q8_match ggml_cuda_match_norm_q8(ggml_cgraph * graph, int i, int device) { + if (graph->nodes[i]->op == GGML_OP_DSV4_HC_POST) { + const auto hc = ggml_cuda_match_hc_post_norm(graph, i); + if (!hc.count || hc.norm->ne[0] % QK8_1 || !ggml_cuda_hc_post_norm_memory_ok(graph, i, hc, device)) { return {}; } + ggml_cuda_norm_q8_match match; + match.post = hc.post; + match.norm = hc.norm; + match.mul = hc.mul; + match.scale = hc.scale; + match.dst = hc.scale ? hc.scale : hc.mul ? hc.mul : hc.norm; + match.last = i + hc.count - 1; + if (!(match.norm->flags & GGML_TENSOR_FLAG_COMPUTE) || !(match.dst->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + return match; + } + ggml_tensor * norm = graph->nodes[i]; + if (norm->op != GGML_OP_RMS_NORM || !norm->src[0] || norm->type != GGML_TYPE_F32 || norm->src[0]->type != GGML_TYPE_F32 || + !(norm->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS}, {}) || + ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE}, {})) { return {}; } + ggml_cuda_norm_q8_match match; + match.norm = match.dst = norm; + match.last = i; + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { + match.mul = graph->nodes[i + 1]; match.add = match.dst = graph->nodes[i + 2]; match.last = i + 2; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { + match.mul = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_SCALE}, {})) { + match.scale = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } + const ggml_tensor * weight = match.mul ? (match.mul->src[0] == norm ? match.mul->src[1] : match.mul->src[0]) : nullptr; + const ggml_tensor * bias = match.add ? (match.add->src[0] == match.mul ? match.add->src[1] : match.add->src[0]) : nullptr; + uintptr_t begin, end; + if (!(match.dst->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_are_same_shape(norm->src[0], norm) || !ggml_are_same_shape(norm, match.dst) || + !ggml_is_contiguous(match.dst) || !ggml_cuda_prepared_range(match.dst, device, begin, end) || + begin % sizeof(float) || (end - begin)/sizeof(float) > INT_MAX) { return {}; } + const ggml_tensor * reads[] = {norm->src[0], weight, bias}; + for (const ggml_tensor * read : reads) { + if (!read) { continue; } + uintptr_t read_begin, read_end; + if (read->type != GGML_TYPE_F32 || read->nb[0] != sizeof(float) || !ggml_cuda_prepared_range(read, device, read_begin, read_end) || + read_begin % sizeof(float) || (read_end - read_begin)/sizeof(float) > INT_MAX || + (begin < read_end && read_begin < end)) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (read->nb[d] % sizeof(float) || read->nb[d]/sizeof(float) > INT64_MAX || read->ne[d] > INT_MAX) { return {}; } + } + } + if (norm->ne[0] % QK8_1 || norm->src[0]->ne[2] > 65535 || norm->src[0]->ne[3] > 65535) { return {}; } + return match; +} + +static std::vector ggml_cuda_plan_norm_q8(ggml_cgraph * graph, int device, + const std::vector & keys, const std::vector & sizes, ggml_cuda_reuse_plan & reuse) { + std::unordered_map> readers; + for (int i = 0; i < graph->n_nodes; ++i) { + if (keys[i] < 0) { continue; } + const ggml_tensor * input = graph->nodes[i]->src[1]; + readers[input->view_src ? input->view_src : input].push_back(i); + } + if (readers.empty()) { return {}; } + std::vector emits(graph->n_nodes); + bool moved = false; + std::vector removed(reuse.groups.size(), false); + for (int i = 0; i < graph->n_nodes; ++i) { + auto match = ggml_cuda_match_norm_q8(graph, i, device); + if (!match.norm) { continue; } + const auto can_emit = [&](int node, int prepare) { + const ggml_tensor * input = graph->nodes[node]->src[1]; + const ggml_tensor * root = input->view_src ? input->view_src : input; + uintptr_t begin, end; + if (prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || + !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } + if (match.post && input->ne[0] % MATRIX_ROW_PADDING && match.norm->ne[0] >= input->ne[0]) { + return false; + } + for (int j = match.last + 1; j <= prepare; ++j) { + if (ggml_cuda_mmvq_input_overwritten(graph->nodes[j], input)) { return false; } + } + return true; + }; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + auto & group = reuse.groups[g]; + if (removed[g] || !can_emit(group.node, group.prepare)) { continue; } + match.image = int(g); + group.prepare = i; + group.after = true; + moved = true; + break; + } + const auto found = readers.find(match.dst); + if (found != readers.end()) { + for (const int node : found->second) { + if (sizes[node] > SIZE_MAX - 255 || !can_emit(node, node)) { continue; } + const int previous = reuse.nodes.empty() ? -1 : reuse.nodes[node]; + if (previous == match.image && match.image >= 0) { continue; } + if (match.image >= 0) { + auto & group = reuse.groups[match.image]; + const ggml_tensor * image = graph->nodes[group.node]->src[1]; + const ggml_tensor * input = graph->nodes[node]->src[1]; + if (input->ne[0] != image->ne[0] || input->ne[1] != image->ne[1]) { + if (input->ne[0] % MATRIX_ROW_PADDING || image->ne[0] % MATRIX_ROW_PADDING || + sizes[node] != sizes[group.node]) { continue; } + } + if (previous >= 0) { + const auto & other = reuse.groups[previous]; + if (!can_emit(other.node, other.prepare)) { continue; } + group.last = std::max(group.last, other.last); + group.size = std::max(group.size, other.size); + for (int & member : reuse.nodes) { + if (member == previous) { member = match.image; } + } + removed[previous] = true; + } else { + reuse.nodes[node] = match.image; + group.last = std::max(group.last, node); + group.size = std::max(group.size, GGML_PAD(sizes[node], 256)); + } + } else { + if (previous >= 0) { continue; } + if (reuse.nodes.empty()) { + reuse.nodes.assign(graph->n_nodes, -1); + reuse.starts.assign(graph->n_nodes, -1); + } + match.image = int(reuse.groups.size()); + reuse.groups.push_back({node, i, node, -1, true, GGML_PAD(sizes[node], 256)}); + removed.push_back(false); + reuse.nodes[node] = match.image; + } + moved = true; + } + } + if (match.image >= 0) { emits[i] = match; } + } + if (moved) { + std::vector remap(reuse.groups.size(), -1); + int next = 0; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + if (removed[g]) { continue; } + remap[g] = next; + reuse.groups[next++] = reuse.groups[g]; + } + reuse.groups.resize(next); + for (int & group : reuse.nodes) { if (group >= 0) { group = remap[group]; } } + for (auto & emit : emits) { if (emit.image >= 0) { emit.image = remap[emit.image]; } } + reuse.pack(); + } + if (reuse.groups.empty()) { return {}; } + return emits; +} + +struct ggml_cuda_norm_mmq_match { + ggml_tensor * post = nullptr; + ggml_tensor * norm = nullptr; + ggml_tensor * mul = nullptr; + ggml_tensor * add = nullptr; + ggml_tensor * scale = nullptr; + ggml_tensor * dst = nullptr; + int last = -1; + int image = -1; +}; + +static ggml_cuda_norm_mmq_match ggml_cuda_match_norm_mmq(ggml_cgraph * graph, int i, int device) { + if (graph->nodes[i]->op == GGML_OP_DSV4_HC_POST) { + const auto match = ggml_cuda_match_norm_emit(graph, i, device); + if (!match.post || match.norm->ne[0] % QK8_1) { return {}; } + ggml_cuda_norm_mmq_match mmq; + mmq.post = match.post; + mmq.norm = match.norm; + mmq.mul = match.mul; + mmq.scale = match.scale; + mmq.dst = match.dst; + mmq.last = match.last; + return mmq; + } + ggml_tensor * norm = graph->nodes[i]; + if (norm->op != GGML_OP_RMS_NORM || !norm->src[0] || norm->type != GGML_TYPE_F32 || norm->src[0]->type != GGML_TYPE_F32 || + !(norm->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS}, {}) || + ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE}, {})) { return {}; } + ggml_cuda_norm_mmq_match match; + match.norm = match.dst = norm; + match.last = i; + if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { + match.mul = graph->nodes[i + 1]; match.add = match.dst = graph->nodes[i + 2]; match.last = i + 2; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { + match.mul = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } else if (ggml_cuda_can_fuse(graph, i, {GGML_OP_RMS_NORM, GGML_OP_SCALE}, {})) { + match.scale = match.dst = graph->nodes[i + 1]; match.last = i + 1; + } + const ggml_tensor * weight = match.mul ? (match.mul->src[0] == norm ? match.mul->src[1] : match.mul->src[0]) : nullptr; + const ggml_tensor * bias = match.add ? (match.add->src[0] == match.mul ? match.add->src[1] : match.add->src[0]) : nullptr; + uintptr_t begin, end; + if (!(match.dst->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_are_same_shape(norm->src[0], norm) || !ggml_are_same_shape(norm, match.dst) || + !ggml_is_contiguous(match.dst) || !ggml_cuda_prepared_range(match.dst, device, begin, end) || + begin % sizeof(float) || (end - begin)/sizeof(float) > INT_MAX) { return {}; } + const ggml_tensor * reads[] = {norm->src[0], weight, bias}; + for (const ggml_tensor * read : reads) { + if (!read) { continue; } + uintptr_t read_begin, read_end; + if (read->type != GGML_TYPE_F32 || read->nb[0] != sizeof(float) || !ggml_cuda_prepared_range(read, device, read_begin, read_end) || + read_begin % sizeof(float) || (read_end - read_begin)/sizeof(float) > INT_MAX || + (begin < read_end && read_begin < end)) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (read->nb[d] % sizeof(float) || read->nb[d]/sizeof(float) > INT64_MAX || read->ne[d] > INT_MAX) { return {}; } + } + } + if (norm->ne[0] % QK8_1 || norm->src[0]->ne[2] > 65535 || norm->src[0]->ne[3] > 65535) { return {}; } + return match; +} + +static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph * graph, int device, + const std::vector & keys, const std::vector & sizes, ggml_cuda_reuse_plan & reuse) { + std::unordered_map> readers; + for (int i = 0; i < graph->n_nodes; ++i) { + if (keys[i] < 0) { continue; } + const ggml_tensor * input = graph->nodes[i]->src[1]; + readers[input->view_src ? input->view_src : input].push_back(i); + } + if (readers.empty()) { return {}; } + std::vector emits(graph->n_nodes); + bool moved = false; + std::vector removed(reuse.groups.size(), false); + for (int i = 0; i < graph->n_nodes; ++i) { + auto match = ggml_cuda_match_norm_mmq(graph, i, device); + if (!match.norm) { continue; } + const auto norm_group_invalid = [&](int node) { + const int group = keys[node] == int(MMQ_Q8_1_DS_LAYOUT_D2S6) ? 64 : 32; + return match.norm->ne[0] % group || graph->nodes[node]->src[1]->ne[0] % group; + }; + const auto can_emit = [&](int node, int prepare) { + const ggml_tensor * input = graph->nodes[node]->src[1]; + const ggml_tensor * root = input->view_src ? input->view_src : input; + uintptr_t begin, end; + if (keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || + norm_group_invalid(node) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || + !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } + for (int j = match.last + 1; j <= prepare; ++j) { + if (ggml_cuda_mmq_input_overwritten(graph->nodes[j], input)) { return false; } + } + return true; + }; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + auto & group = reuse.groups[g]; + if (removed[g] || !can_emit(group.node, group.prepare)) { continue; } + match.image = int(g); + group.prepare = i; + group.after = true; + moved = true; + break; + } + const auto found = readers.find(match.dst); + if (found != readers.end()) { + for (const int node : found->second) { + if (sizes[node] > SIZE_MAX - 255 || !can_emit(node, node)) { continue; } + const int previous = reuse.nodes.empty() ? -1 : reuse.nodes[node]; + if (previous == match.image && match.image >= 0) { continue; } + if (match.image >= 0) { + auto & group = reuse.groups[match.image]; + const ggml_tensor * image = graph->nodes[group.node]->src[1]; + const ggml_tensor * input = graph->nodes[node]->src[1]; + if (keys[node] != keys[group.node] || input->ne[0] != image->ne[0] || input->ne[1] != image->ne[1]) { continue; } + if (previous >= 0) { + const auto & other = reuse.groups[previous]; + if (!can_emit(other.node, other.prepare)) { continue; } + group.last = std::max(group.last, other.last); + group.size = std::max(group.size, other.size); + for (int & member : reuse.nodes) { + if (member == previous) { member = match.image; } + } + removed[previous] = true; + } else { + reuse.nodes[node] = match.image; + group.last = std::max(group.last, node); + group.size = std::max(group.size, GGML_PAD(sizes[node], 256)); + } + } else { + if (previous >= 0) { continue; } + if (reuse.nodes.empty()) { + reuse.nodes.assign(graph->n_nodes, -1); + reuse.starts.assign(graph->n_nodes, -1); + } + match.image = int(reuse.groups.size()); + reuse.groups.push_back({node, i, node, -1, true, GGML_PAD(sizes[node], 256)}); + removed.push_back(false); + reuse.nodes[node] = match.image; + } + moved = true; + } + } + if (match.image >= 0) { emits[i] = match; } + } + if (moved) { + std::vector remap(reuse.groups.size(), -1); + int next = 0; + for (size_t g = 0; g < reuse.groups.size(); ++g) { + if (removed[g]) { continue; } + remap[g] = next; + reuse.groups[next++] = reuse.groups[g]; + } + reuse.groups.resize(next); + for (int & group : reuse.nodes) { if (group >= 0) { group = remap[group]; } } + for (auto & emit : emits) { if (emit.image >= 0) { emit.image = remap[emit.image]; } } + reuse.pack(); + } + if (reuse.groups.empty()) { return {}; } + return emits; +} + + +struct ggml_cuda_scaled_unary_match { + ggml_tensor * first = nullptr; + ggml_tensor * unary = nullptr; + ggml_tensor * last = nullptr; + int count = 0; +}; + +static ggml_cuda_scaled_unary_match ggml_cuda_match_scaled_unary(ggml_cgraph * cgraph, int i) { + ggml_cuda_scaled_unary_match match; + ggml_tensor * first = cgraph->nodes[i]; + const bool before = first->op == GGML_OP_SCALE; + const int u = i + int(before); + if (u >= cgraph->n_nodes || cgraph->nodes[u]->op != GGML_OP_UNARY || + (ggml_get_unary_op(cgraph->nodes[u]) != GGML_UNARY_OP_SILU && ggml_get_unary_op(cgraph->nodes[u]) != GGML_UNARY_OP_SIGMOID)) { + return match; + } + if (before && (ggml_cuda_can_fuse(cgraph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || + ggml_cuda_can_fuse(cgraph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }))) { + return match; + } + if (before && cgraph->nodes[u]->src[0] != first) { + return match; + } + const bool after = u + 1 < cgraph->n_nodes && cgraph->nodes[u + 1]->op == GGML_OP_SCALE && cgraph->nodes[u + 1]->src[0] == cgraph->nodes[u]; + int count = 0; + if (before && after && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { i, i + 1, i + 2 })) { + count = 3; + } else if (before && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY }, { i, i + 1 })) { + count = 2; + } else if (!before && after && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_UNARY, GGML_OP_SCALE }, { i, i + 1 })) { + count = 2; + } + if (!count || !first->src[0] || (first->type != GGML_TYPE_F32 && first->type != GGML_TYPE_BF16)) { + return match; + } + const ggml_tensor * src = first->src[0]; + size_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (src->ne[d] <= 0 || (size_t) src->ne[d] > (INT_MAX - CUDA_NEG_BLOCK_SIZE)/elements) { + return match; + } + elements *= src->ne[d]; + } + for (int j = i - 1; j < i + count; ++j) { + const ggml_tensor * tensor = j < i ? src : cgraph->nodes[j]; + if (tensor->type != first->type || !ggml_is_contiguous(tensor) || !ggml_are_same_shape(src, tensor)) { + return match; + } + } + match.first = first; + match.unary = cgraph->nodes[u]; + match.last = cgraph->nodes[i + count - 1]; + match.count = count; + return match; +} + +static bool ggml_cuda_scaled_unary_memory_ok(const ggml_cuda_scaled_unary_match & match, int device) { + const auto range = [device](const ggml_tensor * tensor, uintptr_t & begin, uintptr_t & end) { + if (!tensor->buffer || !tensor->data || !ggml_backend_buffer_is_cuda(tensor->buffer) || + ((const ggml_backend_cuda_buffer_context *) tensor->buffer->context)->device != device) { + return false; + } + const uintptr_t base = (uintptr_t) ggml_backend_buffer_get_base(tensor->buffer); + const size_t size = ggml_backend_buffer_get_size(tensor->buffer); + const size_t bytes = ggml_nbytes(tensor); + begin = (uintptr_t) tensor->data; + if (begin < base || begin - base > size || bytes > size - (begin - base) || begin > UINTPTR_MAX - bytes || begin % ggml_element_size(tensor)) { + return false; + } + if (tensor->view_src) { + const ggml_tensor * root = tensor->view_src; + if (root->view_src || root->buffer != tensor->buffer || !root->data || tensor->view_offs > ggml_nbytes(root) || + begin != (uintptr_t) root->data + tensor->view_offs || bytes > ggml_nbytes(root) - tensor->view_offs || + (uintptr_t) root->data < base || (uintptr_t) root->data - base > size || ggml_nbytes(root) > size - ((uintptr_t) root->data - base)) { + return false; + } + } + end = begin + bytes; + return true; + }; + const ggml_tensor * tensors[] = { match.first->src[0], match.first, match.unary, match.last }; + uintptr_t begin[4], end[4]; + for (int j = 0; j < 4; ++j) { + if (!range(tensors[j], begin[j], end[j])) { + return false; + } + for (int k = 0; k < j; ++k) { + if (begin[j] < end[k] && begin[k] < end[j] && (begin[j] != begin[k] || end[j] != end[k])) { + return false; + } + } + } + return true; +} + static bool ggml_cuda_hc_injection_bytes(const ggml_tensor * tensor, size_t & bytes) { for (int d = 0; d < GGML_MAX_DIMS; ++d) { if (tensor->ne[d] <= 0) { return false; } @@ -4019,151 +4919,18 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu return match; } -struct ggml_cuda_hc_affine_injection_match { - ggml_cuda_affine_unary_match affine; - ggml_tensor * post = nullptr; - ggml_tensor * comb = nullptr; - ggml_cuda_hc_post_norm_match hc; - int count = 0; -}; - -static bool ggml_cuda_hc_affine_tensor_ok(const ggml_tensor * tensor) { - if (!tensor || tensor->type != GGML_TYPE_F32 || (tensor->view_src && tensor->view_src->type != GGML_TYPE_F32)) { return false; } - size_t elements = 1; - for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (tensor->ne[d] <= 0 || tensor->ne[d] > INT_MAX || tensor->nb[d] % sizeof(float) || elements > INT64_MAX/sizeof(float)/size_t(tensor->ne[d])) { return false; } - elements *= tensor->ne[d]; - } - size_t bytes; - return ggml_cuda_affine_unary_bytes(tensor, bytes) && bytes <= INT64_MAX; -} - -static bool ggml_cuda_hc_affine_prior_source(ggml_cgraph * graph, int start, const ggml_tensor * tensor) { - if (!tensor) { return false; } - if (tensor->view_src) { return ggml_cuda_affine_unary_prior_view(graph, start, tensor); } - if (tensor->op == GGML_OP_NONE) { return true; } - for (int i = 0; i < start; ++i) { - if (graph->nodes[i] == tensor) { return true; } - } - return false; -} - -static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { - auto affine = ggml_cuda_match_affine_unary(ctx, graph, i, allocated); - if (!affine.count || affine.mul->type != GGML_TYPE_F32 || affine.add->type != GGML_TYPE_F32 || affine.unary->type != GGML_TYPE_F32) { return {}; } - ggml_tensor * last = affine.post ? affine.post : affine.unary; - const ggml_tensor * gate_inputs[] = {affine.mul->src[0], affine.mul->src[1], affine.add->src[1]}; - for (const ggml_tensor * tensor : gate_inputs) { - if (!ggml_cuda_hc_affine_tensor_ok(tensor) || tensor->ne[2] != 1 || tensor->ne[3] != 1 || affine.mul->ne[0] % tensor->ne[0] || affine.mul->ne[1] % tensor->ne[1]) { return {}; } - } - int post_i = i + affine.count; - auto skip_views = [&]() { - while (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_VIEW) { - if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[post_i])) { return false; } - ++post_i; - } - return true; - }; - if (!skip_views()) { return {}; } - ggml_tensor * comb = nullptr; - if (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_DSV4_HC_COMB) { - comb = graph->nodes[post_i++]; - // The existing COMB operator supports four streams. - if (comb->type != GGML_TYPE_F32 || comb->ne[0] != 4 || comb->ne[1] != 4 || comb->ne[3] != 1) { return {}; } - const int64_t mix_dim = (comb->ne[0] + 2)*comb->ne[0]; - const ggml_tensor * mixes = comb->src[0]; - const ggml_tensor * scale = comb->src[1]; - const ggml_tensor * base = comb->src[2]; - if (!mixes || !scale || !base || mixes->ne[0] != mix_dim || mixes->ne[1] != comb->ne[2] || mixes->ne[2] != 1 || mixes->ne[3] != 1 || - scale->ne[0] < 3 || scale->ne[1] != 1 || scale->ne[2] != 1 || scale->ne[3] != 1 || - base->ne[0] != mix_dim || base->ne[1] != 1 || base->ne[2] != 1 || base->ne[3] != 1 || comb->op_params[1] <= 0) { return {}; } - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (comb->src[j] && (!ggml_cuda_hc_affine_tensor_ok(comb->src[j]) || !ggml_cuda_hc_affine_prior_source(graph, i, comb->src[j]))) { return {}; } - } - if (!skip_views()) { return {}; } - } - if (post_i >= graph->n_nodes) { return {}; } - ggml_tensor * post = graph->nodes[post_i]; - const ggml_tensor * x = post->src[0]; - const ggml_tensor * residual = post->src[1]; - const ggml_tensor * post_comb = post->src[3]; - if (post->op != GGML_OP_DSV4_HC_POST || post->src[2] != last || !x || !residual || - x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || - !ggml_is_contiguous(post) || !ggml_are_same_shape(post, residual) || x->ne[0] <= 0 || x->ne[1] <= 0 || residual->ne[1] <= 0 || - x->ne[2] != 1 || x->ne[3] != 1 || residual->ne[0] != x->ne[0] || residual->ne[2] != x->ne[1] || residual->ne[3] != 1 || - last->ne[0] != residual->ne[1] || last->ne[1] != x->ne[1] || last->ne[2] != 1 || last->ne[3] != 1 || - (comb && post_comb != comb) || - (post_comb && (post_comb->type != GGML_TYPE_F32 || post_comb->ne[0] != residual->ne[1] || post_comb->ne[1] != residual->ne[1] || post_comb->ne[2] != x->ne[1] || post_comb->ne[3] != 1))) { return {}; } - for (const ggml_tensor * tensor : {x, residual, (const ggml_tensor *) post, post_comb}) { - if (tensor && !ggml_cuda_hc_affine_tensor_ok(tensor)) { return {}; } - } - size_t bytes; - if (!ggml_cuda_hc_injection_bytes(post, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } - auto hc = ggml_cuda_match_hc_post_norm(graph, post_i); - if (hc.count && hc.norm->ne[0] != x->ne[0]) { - if (post_comb && bytes/sizeof(float) <= 256) { return {}; } - if (hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] > WARP_SIZE/2 && hc.norm->ne[0] < WARP_SIZE) { return {}; } - const int block_size = hc.norm->ne[0] < 1024 ? 256 : 1024; - const int64_t nweights = (hc.norm->ne[0] + x->ne[0] - 2)/x->ne[0] + 1; - if ((hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] <= WARP_SIZE/2) || nweights >= block_size) { hc = {}; } - } - const int end = hc.count ? post_i + hc.count : post_i + 1; - const int64_t ncols = hc.count ? hc.norm->ne[0] : 256; - const size_t shared = ((hc.count ? 32 : 0) + (ncols + x->ne[0] - 2)/x->ne[0] + 1)*sizeof(float); - if (shared > ggml_cuda_info().devices[ctx.device].smpb) { return {}; } - std::vector ops; - std::vector outputs; - for (int j = i; j < end; ++j) { - ops.push_back(graph->nodes[j]->op); - if (!hc.mul || graph->nodes[j] != hc.norm) { outputs.push_back(j); } - } - if (!ggml_can_fuse_subgraph(graph, i, end - i, ops.data(), outputs.data(), outputs.size())) { return {}; } - if (allocated) { - if (hc.count && !ggml_cuda_hc_post_norm_memory_ok(graph, post_i, hc, ctx.device)) { return {}; } - std::vector reads = {gate_inputs[0], gate_inputs[1], gate_inputs[2], x, residual, post_comb}; - if (hc.mul) { reads.push_back(hc.mul->src[hc.mul->src[0] == hc.norm ? 1 : 0]); } - if (comb) { - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (comb->src[j]) { reads.push_back(comb->src[j]); } - } - } - std::vector writes = {affine.mul, affine.add, affine.unary}; - if (affine.post) { writes.push_back(affine.post); } - writes.push_back(post); - if (hc.count) { writes.push_back(hc.mul ? hc.mul : hc.norm); } - if (comb) { writes.push_back(comb); } - std::vector rb(reads.size()), re(reads.size()), wb(writes.size()), we(writes.size()); - for (size_t j = 0; j < reads.size(); ++j) { - if (reads[j] && (!ggml_cuda_hc_affine_tensor_ok(reads[j]) || !ggml_cuda_hc_injection_range(reads[j], ctx.device, rb[j], re[j]))) { return {}; } - } - for (size_t j = 0; j < writes.size(); ++j) { - if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_is_contiguous(writes[j]) || !ggml_cuda_hc_affine_tensor_ok(writes[j]) || !ggml_cuda_hc_injection_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } - for (size_t k = 0; k < reads.size(); ++k) { - // COMB is written before the affine kernel reads it. - if (reads[k] && !(writes[j] == comb && reads[k] == comb) && wb[j] < re[k] && rb[k] < we[j]) { return {}; } - } - for (size_t k = 0; k < j; ++k) { - if (wb[j] < we[k] && wb[k] < we[j]) { return {}; } - } - } - for (int j = i; j < end; ++j) { - const ggml_tensor * tensor = graph->nodes[j]; - if (tensor->op != GGML_OP_VIEW && tensor->op != GGML_OP_RESHAPE) { continue; } - uintptr_t begin, stop; - if (!ggml_cuda_hc_injection_range(tensor, ctx.device, begin, stop)) { return {}; } - } - } - return {affine, post, comb, hc, end - i}; -} - -// 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) { +static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, + const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { return 0; } + const auto shared_quantized = [&](const ggml_tensor * mm_node) { + return quantized && mm_node->src[1] == shared_input && ggml_cuda_can_share_mmvq_input(mm_node, cuda_ctx->device) && + !ggml_cuda_mmvq_input_overwritten(mm_node, shared_input) ? quantized : nullptr; + }; ggml_tensor * node = cgraph->nodes[i]; if (node->op == GGML_OP_MUL) { @@ -4494,7 +5261,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data, shared_quantized(up_n)); fused_mul_mat_vec = true; fused_node_count = n_ops; break; @@ -4588,7 +5355,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data, shared_quantized(up_n)); fused_mul_mat_vec = true; fused_node_count = n_ops; break; @@ -4658,7 +5425,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_op = ggml_get_glu_op(glu); fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data, shared_quantized(up_n)); fused_mul_mat_vec = true; fused_node_count = 5; break; @@ -4697,7 +5464,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.glu_op = ggml_get_glu_op(glu); fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3); - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data, shared_quantized(up)); fused_mul_mat_vec = true; fused_node_count = 3; break; @@ -4785,7 +5552,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fusion_data.x_scale = scale; if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data, shared_quantized(mm_node)); fused_mul_mat_vec = true; fused_node_count = n_ops; break; @@ -4850,7 +5617,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data, shared_quantized(mm_node)); fused_mul_mat_vec = true; fused_node_count = 2; break; @@ -4863,7 +5630,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto injection = ggml_cuda_match_hc_injection(*cuda_ctx, cgraph, i); - if (injection.count) { + if (injection.count && (prepared_hc_post < 0 || injection.post != cgraph->nodes[prepared_hc_post] || (!injection.hc.norm && injection.post->src[3]))) { ggml_cuda_op_hc_injection(*cuda_ctx, injection.first, injection.unary, injection.last, injection.post, injection.hc.norm, injection.hc.mul); return injection.count - 1; } @@ -4872,7 +5639,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (node->op == GGML_OP_DSV4_HC_POST && cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto match = ggml_cuda_match_hc_post_norm(cgraph, i); if (match.count != 0 && ggml_cuda_hc_post_norm_memory_ok(cgraph, i, match, cuda_ctx->device)) { - ggml_cuda_op_hc_post_norm(*cuda_ctx, match.post, match.norm, match.mul); + if (match.scale) { + ggml_cuda_op_hc_post_norm_scale(*cuda_ctx, match.post, match.norm, match.scale, nullptr, nullptr); + } else { + ggml_cuda_op_hc_post_norm(*cuda_ctx, match.post, match.norm, match.mul); + } return match.count - 1; } } @@ -4930,15 +5701,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { - const auto injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); - if (injection.count) { - if (injection.comb) { - const bool ok = ggml_cuda_compute_forward(*cuda_ctx, injection.comb); - GGML_ASSERT(ok); - } - const auto & affine = injection.affine; - ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, injection.post, injection.hc.norm, injection.hc.mul); - return injection.count - 1; + const auto match = ggml_cuda_match_scaled_unary(cgraph, i); + if (match.count && ggml_cuda_scaled_unary_memory_ok(match, cuda_ctx->device)) { + ggml_cuda_op_scaled_unary(*cuda_ctx, match.first, match.unary, match.last); + return match.count - 1; + } + } + + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_hc_up(*cuda_ctx, cgraph, i); + if (match.count) { + ggml_cuda_mul_mat_vec_f_hc_pre(*cuda_ctx, match.mm, match.dst); + return match.count - 1; } } @@ -4950,13 +5724,6 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { - const auto match = ggml_cuda_match_hc_up(*cuda_ctx, cgraph, i); - if (match.count) { - ggml_cuda_mul_mat_vec_f_hc_pre(*cuda_ctx, match.mm, match.dst); - return match.count - 1; - } - } return 0; } @@ -4972,24 +5739,6 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_concurrent_event * concurrent_event = nullptr; bool should_launch_concurrent_events = false; - const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { - if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { - concurrent_event = &stream_ctx.concurrent_events[node]; - - is_concurrent_event_active = true; - - GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); - - cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 - GGML_ASSERT(cuda_ctx->curr_stream_no == 0); - CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); - - for (int i = 1; i <= concurrent_event->n_streams; ++i) { - cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); - CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); - } - } - }; while (!graph_evaluated_or_captured) { // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. @@ -5059,6 +5808,146 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud stream_ctx.concurrent_events.clear(); } + static const bool disable_reuse = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); + std::vector input_keys; + std::vector input_sizes; + if (!disable_reuse) { + input_keys.assign(cgraph->n_nodes, -1); + input_sizes.resize(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; ++i) { + input_keys[i] = ggml_cuda_cublas_input_key(cgraph->nodes[i], cuda_ctx->device, input_sizes[i]); + } + } + ggml_cuda_reuse_plan reuse(cgraph, stream_ctx, input_keys, input_sizes, ggml_cuda_prepared_input_overwritten); + std::vector norm_emits; + if (!disable_reuse && stream_ctx.concurrent_events.empty()) { + norm_emits = ggml_cuda_plan_norm_emit(cgraph, cuda_ctx->device, input_keys, input_sizes, reuse); + } + ggml_cuda_reuse_inputs> shared_inputs(cuda_ctx->pool(), reuse); + const auto prepare_group = [&](int group) { + if (shared_inputs[group].get()) { return; } + const auto * node = cgraph->nodes[reuse.groups[group].node]; + ggml_cuda_prepare_cublas_input(*cuda_ctx, node->src[1], input_keys[reuse.groups[group].node], shared_inputs[group]); + }; + const auto prepare_shared = [&](int i, bool after) { + if (reuse.starts.empty()) { return; } + for (int group = reuse.starts[i]; group >= 0; group = reuse.groups[group].next) { + if (reuse.groups[group].after == after) { prepare_group(group); } + } + }; + std::vector mmvq_keys; + std::vector mmvq_sizes; + if (!disable_reuse) { + mmvq_keys.assign(cgraph->n_nodes, -1); + mmvq_sizes.resize(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (ggml_cuda_can_share_mmvq_input(node, cuda_ctx->device)) { + const ggml_tensor * input = node->src[1]; + if (input->ne[0] <= 0 || input->ne[0] > INT64_MAX - MATRIX_ROW_PADDING + 1 || + uint64_t(input->ne[0]) > SIZE_MAX - MATRIX_ROW_PADDING + 1) { continue; } + const size_t blocks = GGML_PAD(size_t(input->ne[0]), MATRIX_ROW_PADDING)/QK8_1; + if (blocks > SIZE_MAX/sizeof(block_q8_1)) { continue; } + const size_t row = blocks*sizeof(block_q8_1); + if (uint64_t(input->ne[1]) > SIZE_MAX/row) { continue; } + mmvq_keys[i] = 0; + mmvq_sizes[i] = size_t(input->ne[1])*row; + } + } + } + ggml_cuda_reuse_plan q8_reuse(cgraph, stream_ctx, mmvq_keys, mmvq_sizes, ggml_cuda_mmvq_input_overwritten); + std::vector q8_emits; + if (!disable_reuse && stream_ctx.concurrent_events.empty()) { + q8_emits = ggml_cuda_plan_norm_q8(cgraph, cuda_ctx->device, mmvq_keys, mmvq_sizes, q8_reuse); + } + ggml_cuda_reuse_inputs> q8_inputs(cuda_ctx->pool(), q8_reuse); + const auto prepare_q8_group = [&](int group) { + if (q8_inputs[group].get()) { return; } + const ggml_tensor * node = cgraph->nodes[q8_reuse.groups[group].node]; + ggml_cuda_quantize_mmvq_input(*cuda_ctx, node->src[0], node->src[1], q8_inputs[group]); + }; + const auto prepare_q8_shared = [&](int i, bool after) { + if (q8_reuse.starts.empty()) { return; } + for (int group = q8_reuse.starts[i]; group >= 0; group = q8_reuse.groups[group].next) { + if (q8_reuse.groups[group].after == after) { + prepare_q8_group(group); + } + } + }; + std::vector mmq_keys; + std::vector mmq_sizes; + const int mmq_cc = ggml_cuda_info().devices[cuda_ctx->device].cc; + const auto scale_size = [&](const ggml_tensor * node) { + return ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) == GGML_PREC_Q4 && node->src[0]->type == GGML_TYPE_NVFP4 ? + GGML_PAD((size_t(node->src[1]->ne[1]) + 128)*sizeof(float), 256) : 0; + }; + if (!disable_reuse) { + mmq_keys.assign(cgraph->n_nodes, -1); + mmq_sizes.resize(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (ggml_cuda_can_share_mmq_input(node, cuda_ctx->device)) { + const ggml_tensor * input = node->src[1]; + if (input->ne[0] <= 0 || input->ne[0] > INT64_MAX - MATRIX_ROW_PADDING + 1 || + uint64_t(input->ne[0]) > SIZE_MAX - MATRIX_ROW_PADDING + 1) { continue; } + const bool fp4 = ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) == GGML_PREC_Q4; + const size_t blocks = GGML_PAD(size_t(input->ne[0]), MATRIX_ROW_PADDING)/(fp4 ? QK_FP4_MMQ : QK8_1_MMQ); + if (blocks > SIZE_MAX/sizeof(block_q8_1_mmq)) { continue; } + const size_t row = blocks*sizeof(block_q8_1_mmq); + // The kernel reads a full tile, including columns past the input. + const size_t tail = 128*sizeof(block_q8_1_mmq); + if (uint64_t(input->ne[1]) > (SIZE_MAX - tail - 255)/row || + uint64_t(input->ne[1]) > (SIZE_MAX - 255)/sizeof(float) - 128) { continue; } + const size_t quantized = size_t(input->ne[1])*row + tail; + if (scale_size(node) > SIZE_MAX - GGML_PAD(quantized, 256)) { continue; } + mmq_keys[i] = ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) == GGML_PREC_Q4 ? + 3 + (node->src[0]->type == GGML_TYPE_NVFP4) : int(mmq_get_q8_1_ds_layout(node->src[0]->type)); + mmq_sizes[i] = GGML_PAD(quantized, 256) + scale_size(node); + } + } + } + ggml_cuda_reuse_plan mmq_reuse(cgraph, stream_ctx, mmq_keys, mmq_sizes, ggml_cuda_mmq_input_overwritten); + std::vector mmq_emits; + if (!disable_reuse && stream_ctx.concurrent_events.empty()) { + mmq_emits = ggml_cuda_plan_norm_mmq(cgraph, cuda_ctx->device, mmq_keys, mmq_sizes, mmq_reuse); + } + ggml_cuda_reuse_inputs mmq_inputs(cuda_ctx->pool(), mmq_reuse); + const auto prepare_mmq_group = [&](int group) { + if (mmq_inputs[group].quantized.get()) { return; } + const ggml_tensor * node = cgraph->nodes[mmq_reuse.groups[group].node]; + ggml_cuda_quantize_mmq_input(*cuda_ctx, node, mmq_reuse.groups[group].size - scale_size(node), mmq_inputs[group]); + }; + const auto prepare_mmq_shared = [&](int i, bool after) { + if (mmq_reuse.starts.empty()) { return; } + for (int group = mmq_reuse.starts[i]; group >= 0; group = mmq_reuse.groups[group].next) { + if (mmq_reuse.groups[group].after == after) { prepare_mmq_group(group); } + } + }; + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { + if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { + const int i = reuse.indices.empty() ? -1 : reuse.indices.at(node); + if (i >= 0) { prepare_shared(i, true); } + const int q8_i = q8_reuse.indices.empty() ? -1 : q8_reuse.indices.at(node); + if (q8_i >= 0) { prepare_q8_shared(q8_i, true); } + const int mmq_i = mmq_reuse.indices.empty() ? -1 : mmq_reuse.indices.at(node); + if (mmq_i >= 0) { prepare_mmq_shared(mmq_i, true); } + concurrent_event = &stream_ctx.concurrent_events[node]; + + is_concurrent_event_active = true; + + GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); + + cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 + GGML_ASSERT(cuda_ctx->curr_stream_no == 0); + CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); + + for (int i = 1; i <= concurrent_event->n_streams; ++i) { + cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); + CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); + } + } + }; + for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -5101,7 +5990,144 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); + const int group = reuse.nodes.empty() ? -1 : reuse.nodes[i]; + prepare_shared(i, false); + if (group >= 0) { prepare_group(group); } + const void * prepared_src1 = group >= 0 ? shared_inputs[group].get() : nullptr; + + prepare_q8_shared(i, false); + const int q8_group = q8_reuse.nodes.empty() ? -1 : q8_reuse.nodes[i]; + if (q8_group >= 0 && !q8_inputs[q8_group].get()) { + bool ready = !is_concurrent_event_active; + if (!ready && !q8_reuse.groups[q8_group].after) { + const auto first = concurrent_event->stream_mapping.find(cgraph->nodes[q8_reuse.groups[q8_group].node]); + ready = first != concurrent_event->stream_mapping.end() && first->second == cuda_ctx->curr_stream_no; + } + if (ready) { prepare_q8_group(q8_group); } + } + const ggml_tensor * shared_input = q8_group >= 0 ? node->src[1] : nullptr; + const char * quantized = q8_group >= 0 ? q8_inputs[q8_group].get() : nullptr; + if (!mmq_emits.empty() && mmq_emits[i].post && (q8_emits.empty() || !q8_emits[i].norm)) { + const auto & emit = mmq_emits[i]; + const auto * typed = !norm_emits.empty() && norm_emits[i].norm ? &norm_emits[i] : nullptr; + if (!typed || (typed->post == emit.post && typed->norm == emit.norm && typed->dst == emit.dst && typed->last == emit.last)) { + const int g = emit.image; + const int reader = mmq_reuse.groups[g].node; + const ggml_tensor * input = cgraph->nodes[reader]->src[1]; + mmq_inputs[g].quantized.alloc(mmq_reuse.groups[g].size); + const auto image = [&](int group) -> void * { + if (group < 0) { return nullptr; } + shared_inputs[group].alloc(input_sizes[reuse.groups[group].node]); + return shared_inputs[group].get(); + }; + void * f16 = typed ? image(typed->f16) : nullptr; + void * bf16 = typed ? image(typed->bf16) : nullptr; + ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + i = emit.last; + continue; + } + } + + if (!q8_emits.empty() && q8_emits[i].norm && + ((!norm_emits.empty() && norm_emits[i].norm) || q8_emits[i].post)) { + const auto & emit = q8_emits[i]; + const int g = emit.image; + const ggml_tensor * input = cgraph->nodes[q8_reuse.groups[g].node]->src[1]; + q8_inputs[g].alloc(mmvq_sizes[q8_reuse.groups[g].node]); + void * f16 = nullptr; + void * bf16 = nullptr; + if (!norm_emits.empty() && norm_emits[i].norm) { + const auto & typed = norm_emits[i]; + GGML_ASSERT(typed.norm == emit.norm && typed.dst == emit.dst && typed.last == emit.last); + const auto image = [&](int group) -> void * { + if (group < 0) { return nullptr; } + shared_inputs[group].alloc(input_sizes[reuse.groups[group].node]); + return shared_inputs[group].get(); + }; + f16 = image(typed.f16); + bf16 = image(typed.bf16); + } + if (emit.post) { + if (emit.scale) { + ggml_cuda_op_dsv4_hc_post(*cuda_ctx, emit.post); + ggml_cuda_op_rms_norm_emit_q8(*cuda_ctx, emit.norm, nullptr, nullptr, emit.scale, f16, bf16, + q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + } else { + ggml_cuda_op_hc_post_norm_emit_q8(*cuda_ctx, emit.post, emit.norm, emit.mul, f16, bf16, + q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + } + } else { + ggml_cuda_op_rms_norm_emit_q8(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, f16, bf16, + q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + } + i = emit.last; + continue; + } + + if (!norm_emits.empty() && norm_emits[i].norm) { + const auto & emit = norm_emits[i]; + const auto image = [&](int g) -> void * { + if (g < 0) { return nullptr; } + shared_inputs[g].alloc(input_sizes[reuse.groups[g].node]); + return shared_inputs[g].get(); + }; + if (emit.post) { + if (emit.scale) { + ggml_cuda_op_hc_post_norm_scale(*cuda_ctx, emit.post, emit.norm, emit.scale, image(emit.f16), image(emit.bf16)); + } else { + ggml_cuda_op_hc_post_norm_emit(*cuda_ctx, emit.post, emit.norm, emit.mul, image(emit.f16), image(emit.bf16)); + } + } else { + ggml_cuda_op_rms_norm_emit(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, image(emit.f16), image(emit.bf16)); + } + i = emit.last; + continue; + } + + if (!q8_emits.empty() && q8_emits[i].norm) { + const auto & emit = q8_emits[i]; + const int g = emit.image; + const ggml_tensor * input = cgraph->nodes[q8_reuse.groups[g].node]->src[1]; + q8_inputs[g].alloc(mmvq_sizes[q8_reuse.groups[g].node]); + ggml_cuda_op_rms_norm_q8(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, + q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + i = emit.last; + continue; + } + + prepare_mmq_shared(i, false); + if (!mmq_emits.empty() && mmq_emits[i].norm) { + const auto & emit = mmq_emits[i]; + const int g = emit.image; + const int reader = mmq_reuse.groups[g].node; + const ggml_tensor * input = cgraph->nodes[reader]->src[1]; + mmq_inputs[g].quantized.alloc(mmq_reuse.groups[g].size); + ggml_cuda_op_rms_norm_mmq(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + i = emit.last; + continue; + } + + const int mmq_group = mmq_reuse.nodes.empty() ? -1 : mmq_reuse.nodes[i]; + if (mmq_group >= 0 && !mmq_inputs[mmq_group].quantized.get()) { + bool ready = !is_concurrent_event_active; + if (!ready && !mmq_reuse.groups[mmq_group].after) { + const auto first = concurrent_event->stream_mapping.find(cgraph->nodes[mmq_reuse.groups[mmq_group].node]); + ready = first != concurrent_event->stream_mapping.end() && first->second == cuda_ctx->curr_stream_no; + } + if (ready) { prepare_mmq_group(mmq_group); } + } + int prepared_hc_post = -1; + if (node->op == GGML_OP_SCALE || node->op == GGML_OP_UNARY) { + for (int j = i + 1; j < cgraph->n_nodes && j - i <= 3; ++j) { + if ((!norm_emits.empty() && norm_emits[j].post) || (!q8_emits.empty() && q8_emits[j].post) || (!mmq_emits.empty() && mmq_emits[j].post)) { + prepared_hc_post = j; + break; + } + } + } + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post); if (nodes_to_skip != 0) { #ifdef GGML_CUDA_DEBUG @@ -5130,7 +6156,16 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud GGML_UNUSED(integrated); #endif // NDEBUG - bool ok = ggml_cuda_compute_forward(*cuda_ctx, node); + bool ok; + if (mmq_group >= 0 && mmq_inputs[mmq_group].quantized.get()) { + ggml_cuda_mul_mat_q(*cuda_ctx, node->src[0], node->src[1], nullptr, node, &mmq_inputs[mmq_group]); + ok = true; + } else if (quantized) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, node->src[0], node->src[1], nullptr, node, nullptr, quantized); + ok = true; + } else { + ok = ggml_cuda_compute_forward(*cuda_ctx, node, prepared_src1); + } if (!ok) { GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); } @@ -5302,22 +6337,12 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { - const auto affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i, false); - if (affine_injection.count) { - ggml_tensor * terminal = cgraph->nodes[i + affine_injection.count - 1]; - ggml_tensor * unused_norm = affine_injection.hc.mul ? affine_injection.hc.norm : nullptr; - for (int j = i; j < i + affine_injection.count - 1; ++j) { - if (cgraph->nodes[j] != unused_norm) { params->add_alloc_dep(params->user_data, cgraph->nodes[j], terminal); } - for (int k = 0; k < GGML_MAX_SRC; ++k) { - if (cgraph->nodes[j]->src[k]) { params->add_alloc_dep(params->user_data, cgraph->nodes[j]->src[k], terminal); } - } - } - for (ggml_tensor * src : terminal->src) { - if (src && src != unused_norm) { params->add_alloc_dep(params->user_data, src, terminal); } - } - i += affine_injection.count - 1; - continue; + const auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); + if (hc_up.count) { + params->add_alloc_dep(params->user_data, hc_up.mm->src[0], hc_up.dst); + params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); if (affine.count) { ggml_tensor * last = cgraph->nodes[i + affine.count - 1]; @@ -5327,11 +6352,6 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph i += affine.count - 1; continue; } - const auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); - if (hc_up.count) { - params->add_alloc_dep(params->user_data, hc_up.mm->src[0], hc_up.dst); - params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); - } const auto injection = ggml_cuda_match_hc_injection(*cuda_ctx, cgraph, i, false); if (injection.count) { diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 3e1721bbe68..4bc02d869d5 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -115,7 +115,7 @@ static ggml_prec ggml_cuda_mmq_get_prec_env() { // src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they // default to the native W4A4 instructions on Blackwell. -static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) { +ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) { static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env(); ggml_prec prec = prec_env; @@ -132,8 +132,63 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm return GGML_PREC_Q4; } +static void ggml_cuda_mmq_clear_padding(const ggml_tensor * src0, cudaStream_t stream) { + // If src0 is a temporary compute buffer, clear any potential padding. + if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size_data = ggml_nbytes(src0); + const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + if (size_alloc > size_data) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + } + } +} + +size_t ggml_cuda_mmq_input_size(const ggml_tensor * node, const int cc) { + const ggml_tensor * src0 = node->src[0]; + const ggml_tensor * src1 = node->src[1]; + const bool use_fp4 = ggml_cuda_mmq_get_prec_src1(src0, node, cc) == GGML_PREC_Q4; + const size_t block_size = use_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq); + const size_t block_values = use_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; + const int J_max = ggml_cuda_mmq_get_J_max(src0->type, src0->ne[1] % 128 != 0, cc, src1->ne[1]); + return src1->ne[3]*src1->ne[2]*src1->ne[1]*GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING)*block_size/block_values + + J_max*sizeof(block_q8_1_mmq); +} + +static void ggml_cuda_quantize_mmq_input_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, + const ggml_tensor * src1, ggml_prec prec, size_t size, ggml_cuda_mmq_input & input) { + input.quantized.alloc(size); + const bool use_fp4 = prec == GGML_PREC_Q4; + if (src0->type == GGML_TYPE_NVFP4 && use_fp4) { + input.scale.alloc(src1->ne[3]*src1->ne[2]*src1->ne[1]); + } + const int64_t padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); + const int64_t s11 = src1->nb[1]/sizeof(float); + const int64_t s12 = src1->nb[2]/sizeof(float); + const int64_t s13 = src1->nb[3]/sizeof(float); + if (use_fp4) { + const bool aligned = ggml_cuda_is_aligned(src1, 32); + quantize_mmq_fp4_cuda((const float *) src1->data, nullptr, input.quantized.get(), input.scale.ptr, src0->type, + aligned, src1->ne[0], s11, s12, s13, padded, src1->ne[1], src1->ne[2], src1->ne[3], ctx.stream()); + } else { + quantize_mmq_q8_1_cuda((const float *) src1->data, nullptr, input.quantized.get(), src0->type, + src1->ne[0], s11, s12, s13, padded, src1->ne[1], src1->ne[2], src1->ne[3], ctx.stream()); + } + CUDA_CHECK(cudaGetLastError()); +} + +void ggml_cuda_quantize_mmq_input(ggml_backend_cuda_context & ctx, const ggml_tensor * node, + size_t size, ggml_cuda_mmq_input & input) { + const int cc = ggml_cuda_info().devices[ctx.device].cc; + ggml_cuda_quantize_mmq_input_impl(ctx, node->src[0], node->src[1], + ggml_cuda_mmq_get_prec_src1(node->src[0], node, cc), size, input); +} + void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_cuda_mmq_input * input) { + GGML_ASSERT(!ids || !input); GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. @@ -156,16 +211,7 @@ void ggml_cuda_mul_mat_q( const float * src1_d = (const float *) src1->data; float * dst_d = (float *) dst->data; - // If src0 is a temporary compute buffer, clear any potential padding. - if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { - const size_t size_data = ggml_nbytes(src0); - const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); - if (size_alloc > size_data) { - GGML_ASSERT(ggml_is_contiguously_allocated(src0)); - GGML_ASSERT(!src0->view_src); - CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); - } - } + ggml_cuda_mmq_clear_padding(src0, stream); const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); @@ -185,30 +231,13 @@ void ggml_cuda_mul_mat_q( const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ : QK8_1_MMQ; if (!ids) { - const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + - ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq); - ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); - ggml_cuda_pool_alloc src1_scale(ctx.pool()); - if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) { - src1_scale.alloc(ne13*ne12*ne11); - } - - { - const int64_t s11 = src1->nb[1] / ts_src1; - const int64_t s12 = src1->nb[2] / ts_src1; - const int64_t s13 = src1->nb[3] / ts_src1; - if (use_native_fp4) { - static constexpr size_t align_float8 = 32; - const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8); - static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); - quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, - ne11, ne12, ne13, stream); - - } else { - quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, - ne11, ne12, ne13, stream); - } - CUDA_CHECK(cudaGetLastError()); + ggml_cuda_mmq_input local_input(ctx.pool()); + if (!input) { + const int J_max = ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11); + const size_t size = ne13*ne12*ne11*ne10_padded*y_block_size/y_values_per_block + + J_max*sizeof(block_q8_1_mmq); + ggml_cuda_quantize_mmq_input_impl(ctx, src0, src1, prec_src1, size, local_input); + input = &local_input; } // Stride depends on quantization format @@ -218,8 +247,8 @@ void ggml_cuda_mul_mat_q( const int64_t s13 = ne12*s12; const mmq_args args = { - src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d, - src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr, + src0_d, src0->type, (const int *) input->quantized.ptr, nullptr, nullptr, dst_d, + src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? input->scale.ptr : nullptr, ne00, ne01, ne1, s01, ne11, s1, ne02, ne12, s02, s12, s2, ne03, ne13, s03, s13, s3, diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 0cd31d91648..1ea6b427d5c 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -1613,7 +1613,22 @@ extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4); // ------------------------------------------------------------------------------------------------------------------------- +struct ggml_cuda_mmq_input { + ggml_cuda_pool_alloc quantized; + ggml_cuda_pool_alloc scale; + + explicit ggml_cuda_mmq_input(ggml_cuda_pool & pool) : quantized(pool), scale(pool) {} +}; + +size_t ggml_cuda_mmq_input_size(const ggml_tensor * node, int cc); + +void ggml_cuda_quantize_mmq_input(ggml_backend_cuda_context & ctx, const ggml_tensor * node, + size_t size, ggml_cuda_mmq_input & input); + void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_cuda_mmq_input * input = nullptr); bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts); + +ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, int cc); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index a35dc937073..9581632d9c2 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1418,9 +1418,18 @@ static void mul_mat_vec_q_switch_type( } } +void ggml_cuda_quantize_mmvq_input(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, + const ggml_tensor * src1, ggml_cuda_pool_alloc & quantized) { + const int64_t padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); + quantized.alloc(src1->ne[3]*src1->ne[2]*src1->ne[1]*padded*sizeof(block_q8_1)/QK8_1); + quantize_row_q8_1_cuda((const float *) src1->data, nullptr, quantized.get(), src0->type, src1->ne[0], + src1->nb[1]/sizeof(float), src1->nb[2]/sizeof(float), src1->nb[3]/sizeof(float), + padded, src1->ne[1], src1->ne[2], src1->ne[3], ctx.stream()); +} + void ggml_cuda_mul_mat_vec_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, - const ggml_cuda_mm_fusion_args_host * fusion) { + const ggml_cuda_mm_fusion_args_host * fusion, const char * quantized) { GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. @@ -1440,7 +1449,6 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE); - const float * src1_d = (const float *) src1->data; const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; float * dst_d = (float *) dst->data; @@ -1498,12 +1506,10 @@ void ggml_cuda_mul_mat_vec_q( } const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); - ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); - { - const int64_t s11 = src1->nb[1] / ts_src1; - const int64_t s12 = src1->nb[2] / ts_src1; - const int64_t s13 = src1->nb[3] / ts_src1; - quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); + if (!quantized) { + ggml_cuda_quantize_mmvq_input(ctx, src0, src1, src1_q8_1); + quantized = src1_q8_1.get(); } const int64_t s01 = src0->nb[1] / ts_src0; @@ -1529,7 +1535,7 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; mul_mat_vec_q_switch_type( - src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, + src0->data, src0->type, quantized, ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, ne03, ne3, s03, s13, s3, ids_stride, stream); diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e6..c3c1597b501 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -9,7 +9,10 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); int get_mmvq_mmid_max_batch(ggml_type type, int cc); void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr); + const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr, const char * quantized = nullptr); + +void ggml_cuda_quantize_mmvq_input(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, + const ggml_tensor * src1, ggml_cuda_pool_alloc & quantized); void ggml_cuda_op_mul_mat_vec_q( ggml_backend_cuda_context & ctx, diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 27d3c4c0f67..355fc21e493 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1,4 +1,7 @@ #include "norm.cuh" +#include "mmq.cuh" +#include "quantize.cuh" +#include "convert.cuh" #include "unary.cuh" #include @@ -74,14 +77,251 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr } } -template -static __global__ void rms_norm_f32(const float * x, +struct ggml_cuda_norm_store { + static constexpr int width = 1; + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + dst[col] = value; + GGML_UNUSED(base); + } +}; + +struct ggml_cuda_norm_emit_store { + static constexpr int width = 1; + half * f16; + nv_bfloat16 * bf16; + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + dst[col] = value; + const int64_t index = dst - base + col; + if (f16) { f16[index] = ggml_cuda_cast(value); } + if (bf16) { bf16[index] = ggml_cuda_cast(value); } + } +}; + +struct ggml_cuda_norm_q8_store { + static constexpr int width = 1; + block_q8_1 * image; + int64_t cols; + int64_t padded; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + dst[col] = value; + const int64_t index = dst - base + col; + const int64_t row = index / cols; + const int64_t column = index % cols; + const int lane = column % QK8_1; + const int64_t block = (row*padded + column)/QK8_1; + const float amax = warp_reduce_max(fabsf(value)); + const float sum = warp_reduce_sum(value); + const float d = amax/127.0f; + image[block].qs[lane] = amax == 0.0f ? 0 : roundf(value/d); + if (lane == 0) { image[block].ds = make_half2(d, sum); } + if (column/QK8_1 == cols/QK8_1 - 1) { + for (int64_t tail = cols/QK8_1; tail < padded/QK8_1; ++tail) { + block_q8_1 & zero = image[row*(padded/QK8_1) + tail]; + zero.qs[lane] = 0; + if (lane == 0) { zero.ds = make_half2(0.0f, 0.0f); } + } + } + } +}; + +struct ggml_cuda_norm_emit_q8_store { + static constexpr int width = 1; + block_q8_1 * image; + uint3 cols; + int64_t padded; + half * f16; + nv_bfloat16 * bf16; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + dst[col] = value; + const int64_t index = dst - base + col; + const uint2 rc = fast_div_modulo(uint32_t(index), cols); + const int64_t row = rc.x; + const int64_t column = rc.y; + const int lane = column % QK8_1; + const int64_t block = (row*padded + column)/QK8_1; + const float amax = warp_reduce_max(fabsf(value)); + const float sum = warp_reduce_sum(value); + const float d = amax/127.0f; + image[block].qs[lane] = amax == 0.0f ? 0 : roundf(value/d); + if (lane == 0) { image[block].ds = make_half2(d, sum); } + if (column/QK8_1 == cols.z/QK8_1 - 1) { + for (int64_t tail = cols.z/QK8_1; tail < padded/QK8_1; ++tail) { + block_q8_1 & zero = image[row*(padded/QK8_1) + tail]; + zero.qs[lane] = 0; + if (lane == 0) { zero.ds = make_half2(0.0f, 0.0f); } + } + } + if (f16) { f16[index] = ggml_cuda_cast(value); } + if (bf16) { bf16[index] = ggml_cuda_cast(value); } + } +}; + +template +struct ggml_cuda_norm_mmq_store { + static constexpr int width = 4; + block_q8_1_mmq * image; + int64_t cols; + int64_t padded; + int64_t rows; + + __device__ __forceinline__ void emit(int64_t index, float4 xi) const { + constexpr int vals_per_scale = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + const int64_t column = index % cols; + const int64_t row = index / cols; + store(column, row, xi); + if (column >= cols - vals_per_scale) { + for (int64_t tail = cols; tail < padded; tail += vals_per_scale) { + store(tail + column % vals_per_scale, row, make_float4(0.0f, 0.0f, 0.0f, 0.0f)); + } + } + } + + __device__ __forceinline__ void store(int64_t column, int64_t row, float4 xi) const { + constexpr int vals_per_scale = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + constexpr int vals_per_sum = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; + const int64_t ib = (column/QK8_1_MMQ)*rows + row; + const int iqs = column % QK8_1_MMQ; + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); + amax = fmaxf(amax, fabsf(xi.w)); + const unsigned mask = __activemask(); +#pragma unroll + for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, offset, WARP_SIZE)); + } + float sum; + if constexpr (Layout != MMQ_Q8_1_DS_LAYOUT_D4) { + sum = xi.x + xi.y + xi.z + xi.w; +#pragma unroll + for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) { + sum += __shfl_xor_sync(mask, sum, offset, WARP_SIZE); + } + } + const float d_inv = 127.0f/amax; + const float d = 1.0f/d_inv; + 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); + ((char4 *) image[ib].qs)[iqs/4] = q; + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { + if (iqs % 16 == 0 && iqs < 96) { + image[ib].d2s6[2 + iqs/16] = sum; + if (iqs % 64 == 0) { image[ib].d2s6[iqs/64] = d; } + } + } else if (iqs % 32 == 0) { + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_DS4) { image[ib].ds4[iqs/32] = make_half2(d, sum); } + else { image[ib].d4[iqs/32] = d; } + } + } + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float4 value) const { + dst[col] = value.x; dst[col + 1] = value.y; + dst[col + 2] = value.z; dst[col + 3] = value.w; + emit(dst - base + col, value); + } +}; + +struct ggml_cuda_norm_mxfp4_store { + static constexpr int width = 4; + block_fp4_mmq * image; + int64_t cols; + int64_t padded; + int64_t rows; + + __device__ __forceinline__ void store(int64_t column, int64_t row, float4 value) const { + const int lane = (column % 32) / 4; + float amax = fabsf(value.x); + amax = fmaxf(amax, fabsf(value.y)); + amax = fmaxf(amax, fabsf(value.z)); + amax = fmaxf(amax, fabsf(value.w)); + const unsigned mask = __activemask(); +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, offset, 8)); + } + uint8_t e = 0; + if (amax > 0.0f) { + e = static_cast(min(max(__float2int_rn(log2f(amax)) - 2 + 127, 0), 254)); + } + const float inv_s = amax == 0.0f ? 0.0f : __frcp_rn(ggml_cuda_e8m0_to_fp32(e)); + const int sender = lane / 2; +#if CUDART_VERSION >= 12080 + value.x *= inv_s; value.y *= inv_s; value.z *= inv_s; value.w *= inv_s; +#else + value.x = ggml_cuda_float_to_fp4_e2m1(value.x, inv_s); + value.y = ggml_cuda_float_to_fp4_e2m1(value.y, inv_s); + value.z = ggml_cuda_float_to_fp4_e2m1(value.z, inv_s); + value.w = ggml_cuda_float_to_fp4_e2m1(value.w, inv_s); +#endif + const float lo0 = __shfl_sync(mask, value.x, sender, 8); + const float lo1 = __shfl_sync(mask, value.y, sender, 8); + const float lo2 = __shfl_sync(mask, value.z, sender, 8); + const float lo3 = __shfl_sync(mask, value.w, sender, 8); + const float hi0 = __shfl_sync(mask, value.x, sender + 4, 8); + const float hi1 = __shfl_sync(mask, value.y, sender + 4, 8); + const float hi2 = __shfl_sync(mask, value.z, sender + 4, 8); + const float hi3 = __shfl_sync(mask, value.w, sender + 4, 8); + const float v0 = lane % 2 ? lo2 : lo0; + const float v1 = lane % 2 ? hi2 : hi0; + const float v2 = lane % 2 ? lo3 : lo1; + const float v3 = lane % 2 ? hi3 : hi1; + char2 packed; +#if CUDART_VERSION >= 12080 + const __nv_fp4x4_e2m1 fp4(make_float4(v0, v1, v2, v3)); + packed = *reinterpret_cast(&fp4); +#else + packed = make_char2((uint8_t(v1) << 4) | uint8_t(v0), (uint8_t(v3) << 4) | uint8_t(v2)); +#endif + const int sub = (column % QK_FP4_MMQ) / 32; + block_fp4_mmq & block = image[(column / QK_FP4_MMQ) * rows + row]; + reinterpret_cast(block.qs)[sub * 8 + lane] = packed; + if (lane == 0) { + uint8_t * header = reinterpret_cast(block.d4) + (sub / 2) * sizeof(uint32_t); + header[sub % 2] = e; + if (sub % 2 == 0) { header[2] = 0; header[3] = 0; } + } + } + + __device__ __forceinline__ void zero(int64_t column, int64_t row) const { + const int lane = (column % 32) / 4; + const int sub = (column % QK_FP4_MMQ) / 32; + block_fp4_mmq & block = image[(column / QK_FP4_MMQ) * rows + row]; + reinterpret_cast(block.qs)[sub * 8 + lane] = make_char2(0, 0); + if (lane == 0) { + uint8_t * header = reinterpret_cast(block.d4) + (sub / 2) * sizeof(uint32_t); + header[sub % 2] = 0; + if (sub % 2 == 0) { header[2] = 0; header[3] = 0; } + } + } + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float4 value) const { + dst[col] = value.x; dst[col + 1] = value.y; dst[col + 2] = value.z; dst[col + 3] = value.w; + const int64_t index = dst - base + col; + const int64_t row = index / cols; + const int64_t column = index % cols; + store(column, row, value); + if (column >= cols - 32) { + for (int64_t tail = cols; tail < padded; tail += 32) { + zero(tail + column % 32, row); + } + } + } +}; + +template +static __device__ __forceinline__ void rms_norm_f32_impl(const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, + const Write write, const float * mul = nullptr, const int64_t mul_stride_row = 0, const int64_t mul_stride_channel = 0, @@ -111,6 +351,7 @@ static __global__ void rms_norm_f32(const float * x, static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying"); + const float * dst_base = dst; x += sample*stride_sample + channel*stride_channel + row*stride_row; dst += ((sample*nchannels + channel)*nrows + row)*ncols; @@ -143,20 +384,233 @@ static __global__ void rms_norm_f32(const float * x, const float mean = tmp / ncols; const float scale = rsqrtf(mean + eps); + if constexpr (Write::width == 4) { + for (int col = 4*tid; col < ncols; col += 4*block_size) { + float values[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int column = col + j; + if constexpr (do_multiply && do_add) { + const int mul_col = fastmodulo(column, mul_ncols_packed); + const int add_col = fastmodulo(column, add_ncols_packed); + values[j] = scale * x[column] * mul[mul_col] + add[add_col]; + } else if constexpr (do_multiply) { + const int mul_col = fastmodulo(column, mul_ncols_packed); + values[j] = scale * x[column] * mul[mul_col]; + } else if constexpr (do_scale) { + values[j] = scale_out * (scale * x[column]); + } else { values[j] = scale * x[column]; } + } + write(dst, dst_base, col, make_float4(values[0], values[1], values[2], values[3])); + } + } else { for (int col = tid; col < ncols; col += block_size) { if constexpr (do_multiply && do_add) { const int mul_col = fastmodulo(col, mul_ncols_packed); const int add_col = fastmodulo(col, add_ncols_packed); - dst[col] = scale * x[col] * mul[mul_col] + add[add_col]; + write(dst, dst_base, col, scale * x[col] * mul[mul_col] + add[add_col]); } else if constexpr (do_multiply) { const int mul_col = fastmodulo(col, mul_ncols_packed); - dst[col] = scale * x[col] * mul[mul_col]; + write(dst, dst_base, col, scale * x[col] * mul[mul_col]); } else if constexpr (do_scale) { - dst[col] = scale_out * (scale * x[col]); + write(dst, dst_base, col, scale_out * (scale * x[col])); } else { - dst[col] = scale * x[col]; + write(dst, dst_base, col, scale * x[col]); } } + } +} + +template +static __global__ void rms_norm_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_store{}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + +template +static __global__ void rms_norm_emit_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + half * f16, + nv_bfloat16 * bf16, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_emit_store{f16, bf16}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + +template +static __global__ void rms_norm_q8_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + void * image, + int64_t cols, + int64_t padded, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_q8_store{(block_q8_1 *) image, cols, padded}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + +template +static __global__ void rms_norm_emit_q8_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + half * f16, + nv_bfloat16 * bf16, + void * image, + uint3 cols, + int64_t padded, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_emit_q8_store{(block_q8_1 *) image, cols, padded, f16, bf16}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + +template +static __global__ void rms_norm_mmq_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + void * image, + int64_t cols, + int64_t padded, + int64_t rows, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_mmq_store{(block_q8_1_mmq *) image, cols, padded, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + +template +static __global__ void rms_norm_mxfp4_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + void * image, + int64_t cols, + int64_t padded, + int64_t rows, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_mxfp4_store{(block_fp4_mmq *) image, cols, padded, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); } template @@ -731,6 +1185,126 @@ void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { l2_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +template +static void ggml_cuda_rms_norm_emit_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + ggml_cuda_kernel_launch(rms_norm_emit_f32, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), (half *) f16, (nv_bfloat16 *) bf16, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_emit(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_emit_launch<256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, f16, bf16); + } else { + ggml_cuda_rms_norm_emit_launch<1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, f16, bf16); + } +} + +void ggml_cuda_op_rms_norm_emit(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_emit(ctx, norm, mul, add, scale, f16, bf16); } + else if (mul) { ggml_cuda_rms_norm_emit(ctx, norm, mul, add, scale, f16, bf16); } + else if (scale) { ggml_cuda_rms_norm_emit(ctx, norm, mul, add, scale, f16, bf16); } + else { ggml_cuda_rms_norm_emit(ctx, norm, mul, add, scale, f16, bf16); } +} + +template +struct ggml_cuda_hc_norm_emit_mmq_pair_store { + static constexpr int width = 2; + block_q8_1_mmq * image; + uint3 cols; + int64_t padded; + int64_t rows; + half * f16; + nv_bfloat16 * bf16; + float scale; + + __device__ __forceinline__ void store(int64_t column, int64_t row, float2 xi) const { + constexpr int vals_per_scale = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + constexpr int vals_per_sum = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; + const unsigned mask = __activemask(); + float amax = fmaxf(fabsf(xi.x), fabsf(xi.y)); +#pragma unroll + for (int offset = vals_per_scale/4; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(mask, amax, offset, WARP_SIZE)); + } + float sum; + if constexpr (Layout != MMQ_Q8_1_DS_LAYOUT_D4) { + const int pair = threadIdx.x % WARP_SIZE & ~1; + const float x = __shfl_sync(mask, xi.x, pair, WARP_SIZE); + const float y = __shfl_sync(mask, xi.y, pair, WARP_SIZE); + const float z = __shfl_sync(mask, xi.x, pair + 1, WARP_SIZE); + const float w = __shfl_sync(mask, xi.y, pair + 1, WARP_SIZE); + sum = x + y + z + w; +#pragma unroll + for (int offset = vals_per_sum/4; offset >= 2; offset >>= 1) { + sum += __shfl_xor_sync(mask, sum, offset, WARP_SIZE); + } + } + const float d_inv = 127.0f/amax; + const float d = 1.0f/d_inv; + const int64_t ib = (column/QK8_1_MMQ)*rows + row; + const int iqs = column % QK8_1_MMQ; + char2 q; + q.x = roundf(xi.x*d_inv); + q.y = roundf(xi.y*d_inv); + ((char2 *) image[ib].qs)[iqs/2] = q; + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { + if (iqs % 16 == 0 && iqs < 96) { + image[ib].d2s6[2 + iqs/16] = sum; + if (iqs % 64 == 0) { image[ib].d2s6[iqs/64] = d; } + } + } else if (iqs % 32 == 0) { + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_DS4) { image[ib].ds4[iqs/32] = make_half2(d, sum); } + else { image[ib].d4[iqs/32] = d; } + } + } + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float2 xi) const { + if constexpr (Scale) { + xi.x = scale * xi.x; + xi.y = scale * xi.y; + } + dst[col] = xi.x; + dst[col + 1] = xi.y; + const int64_t index = dst - base + col; + const uint2 rc = fast_div_modulo(uint32_t(index), cols); + store(rc.y, rc.x, xi); + constexpr int vals_per_scale = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + if (rc.y >= cols.z - vals_per_scale) { + for (int64_t tail = cols.z; tail < padded; tail += vals_per_scale) { + store(tail + rc.y % vals_per_scale, rc.x, make_float2(0.0f, 0.0f)); + } + } + if (f16) { + f16[index] = ggml_cuda_cast(xi.x); + f16[index + 1] = ggml_cuda_cast(xi.y); + } + if (bf16) { + bf16[index] = ggml_cuda_cast(xi.x); + bf16[index + 1] = ggml_cuda_cast(xi.y); + } + } +}; + struct hc_post_norm_data { const float * x; const float * residual; @@ -797,24 +1371,154 @@ static __global__ __launch_bounds__(block_size) void hc_post_norm_f32(hc_post_no } } -template -static auto hc_post_norm_kernel(int layout) { - if (layout == 1) { - return hc_post_norm_f32; - } - if (layout == 2) { - return hc_post_norm_f32; +template +static __device__ __forceinline__ void hc_post_norm_f32_impl(hc_post_norm_data a, const Write write) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; } - return hc_post_norm_f32; -} - -bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * mul) { - const ggml_tensor * dst = mul ? mul : norm; - if (!post->src[0] || !post->src[1] || !post->src[2] || post->op != GGML_OP_DSV4_HC_POST || norm->op != GGML_OP_RMS_NORM || - post->type != GGML_TYPE_F32 || norm->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || - !ggml_is_contiguous(post) || !ggml_is_contiguous(norm) || !ggml_is_contiguous(dst) || - norm->ne[0] > INT_MAX - 1024 || norm->ne[1] > INT_MAX || norm->ne[2] > 65535 || norm->ne[3] > 65535) { - return false; + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + ggml_cuda_pdl_sync(); + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * a.post[ih*a.sp[0] + it*a.sp[1]]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + extern __shared__ float s_sum[]; + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + if constexpr (Write::width == 2) { + for (int64_t col = 2*threadIdx.x; col < a.ncols; col += 2*block_size) { + float x, y; + if constexpr (do_multiply) { + x = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + y = scale * a.post_dst[offset + col + 1] * a.mul[fastmodulo(col + 1, a.mul_ne[0])]; + } else { + x = scale * a.post_dst[offset + col]; + y = scale * a.post_dst[offset + col + 1]; + } + write(a.dst + offset, a.dst, int(col), make_float2(x, y)); + } + } else { + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + write(a.dst + offset, a.dst, col, scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]); + } else { + write(a.dst + offset, a.dst, col, scale * a.post_dst[offset + col]); + } + } + } +} + +template +static __global__ __launch_bounds__(block_size) void hc_post_norm_emit_f32(hc_post_norm_data a, half * f16, nv_bfloat16 * bf16) { + hc_post_norm_f32_impl(a, ggml_cuda_norm_emit_store{f16, bf16}); +} + +template +static __global__ __launch_bounds__(block_size) void hc_post_norm_emit_q8_f32(hc_post_norm_data a, half * f16, nv_bfloat16 * bf16, void * image, uint3 cols, int64_t padded) { + hc_post_norm_f32_impl(a, ggml_cuda_norm_emit_q8_store{(block_q8_1 *) image, cols, padded, f16, bf16}); +} + +template +static __global__ __launch_bounds__(Block) void hc_post_norm_emit_mmq_f32(hc_post_norm_data a, + half * f16, nv_bfloat16 * bf16, void * image, uint3 cols, int64_t padded, int64_t rows, float scale) { + hc_post_norm_f32_impl(a, + ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) image, cols, padded, rows, f16, bf16, scale}); +} + +template +static auto hc_post_norm_emit_mmq_kernel(int layout) { + if (layout == 1) { return hc_post_norm_emit_mmq_f32; } + if (layout == 2) { return hc_post_norm_emit_mmq_f32; } + return hc_post_norm_emit_mmq_f32; +} + +struct ggml_cuda_norm_scale_store { + static constexpr int width = 1; + float scale; + half * f16; + nv_bfloat16 * bf16; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + value = scale * value; + ggml_cuda_norm_emit_store{f16, bf16}(dst, base, col, value); + } +}; + +template +static __global__ __launch_bounds__(Block) void hc_post_norm_scale_f32(hc_post_norm_data a, float scale, + half * f16, nv_bfloat16 * bf16) { + hc_post_norm_f32_impl(a, ggml_cuda_norm_scale_store{scale, f16, bf16}); +} + +template +static auto hc_post_norm_scale_kernel(int layout) { + if (layout == 1) { return hc_post_norm_scale_f32; } + if (layout == 2) { return hc_post_norm_scale_f32; } + return hc_post_norm_scale_f32; +} + +template +static auto hc_post_norm_kernel(int layout) { + if (layout == 1) { + return hc_post_norm_f32; + } + if (layout == 2) { + return hc_post_norm_f32; + } + return hc_post_norm_f32; +} + +template +static auto hc_post_norm_emit_kernel(int layout) { + if (layout == 1) { + return hc_post_norm_emit_f32; + } + if (layout == 2) { + return hc_post_norm_emit_f32; + } + return hc_post_norm_emit_f32; +} + +template +static auto hc_post_norm_emit_q8_kernel(int layout) { + if (layout == 1) { + return hc_post_norm_emit_q8_f32; + } + if (layout == 2) { + return hc_post_norm_emit_q8_f32; + } + return hc_post_norm_emit_q8_f32; +} + +bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * mul) { + const ggml_tensor * dst = mul ? mul : norm; + if (!post->src[0] || !post->src[1] || !post->src[2] || post->op != GGML_OP_DSV4_HC_POST || norm->op != GGML_OP_RMS_NORM || + post->type != GGML_TYPE_F32 || norm->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + !ggml_is_contiguous(post) || !ggml_is_contiguous(norm) || !ggml_is_contiguous(dst) || + norm->ne[0] > INT_MAX - 1024 || norm->ne[1] > INT_MAX || norm->ne[2] > 65535 || norm->ne[3] > 65535) { + return false; } const ggml_tensor * inputs[] = { post->src[0], post->src[1], post->src[2], post->src[3], post, norm, dst, mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr }; for (const ggml_tensor * tensor : inputs) { @@ -843,7 +1547,164 @@ bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_ten return ggml_nelements(norm) <= INT_MAX && ggml_nelements(post) == ggml_nelements(norm) && ggml_are_same_shape(norm, dst) && std::isfinite(eps) && eps >= 0.0f; } +bool ggml_cuda_should_fuse_hc_post_norm_scale(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * scale) { + return scale && ggml_cuda_should_fuse_hc_post_norm(post, norm, nullptr) && scale->op == GGML_OP_SCALE && + scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_are_same_shape(norm, scale) && + (norm->flags & GGML_TENSOR_FLAG_COMPUTE) && (scale->flags & GGML_TENSOR_FLAG_COMPUTE); +} + +template +static void ggml_cuda_op_hc_post_norm_impl(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, void * f16, void * bf16, + ggml_tensor * scale = nullptr) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = (float *) (postop ? scale : (mul ? mul : norm))->data; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm->ne[0]; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = ggml_get_op_params_f32(norm, 0); + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + if constexpr (postop) { + auto kernel = block_size == 256 + ? (comb ? hc_post_norm_scale_kernel<256, true>(layout) : hc_post_norm_scale_kernel<256, false>(layout)) + : (comb ? hc_post_norm_scale_kernel<1024, true>(layout) : hc_post_norm_scale_kernel<1024, false>(layout)); + ggml_cuda_kernel_launch(kernel, launch, a, ggml_get_op_params_f32(scale, 0), (half *) f16, (nv_bfloat16 *) bf16); + } else if constexpr (emit) { + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_post_norm_emit_kernel<256, true, true>(layout) : hc_post_norm_emit_kernel<256, true, false>(layout)) : (mul ? hc_post_norm_emit_kernel<256, false, true>(layout) : hc_post_norm_emit_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_post_norm_emit_kernel<1024, true, true>(layout) : hc_post_norm_emit_kernel<1024, true, false>(layout)) : (mul ? hc_post_norm_emit_kernel<1024, false, true>(layout) : hc_post_norm_emit_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16); + } else { + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_post_norm_kernel<256, true, true>(layout) : hc_post_norm_kernel<256, true, false>(layout)) : (mul ? hc_post_norm_kernel<256, false, true>(layout) : hc_post_norm_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_post_norm_kernel<1024, true, true>(layout) : hc_post_norm_kernel<1024, true, false>(layout)) : (mul ? hc_post_norm_kernel<1024, false, true>(layout) : hc_post_norm_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a); + } +} + void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { + ggml_cuda_op_hc_post_norm_impl(ctx, post, norm, mul, nullptr, nullptr); +} + +void ggml_cuda_op_hc_post_norm_emit(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, void * f16, void * bf16) { + ggml_cuda_op_hc_post_norm_impl(ctx, post, norm, mul, f16, bf16); +} + +void ggml_cuda_op_hc_post_norm_scale(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, + ggml_tensor * scale, void * f16, void * bf16) { + GGML_ASSERT(ggml_cuda_should_fuse_hc_post_norm_scale(post, norm, scale)); + ggml_cuda_op_hc_post_norm_impl(ctx, post, norm, nullptr, f16, bf16, scale); +} + +template +static void ggml_cuda_rms_norm_q8_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + ggml_cuda_kernel_launch(rms_norm_q8_f32, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), image, cols, padded, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_q8_launch<256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded); + } else { + ggml_cuda_rms_norm_q8_launch<1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded); + } +} + +void ggml_cuda_op_rms_norm_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_q8(ctx, norm, mul, add, scale, image, cols, padded); } + else if (mul) { ggml_cuda_rms_norm_q8(ctx, norm, mul, add, scale, image, cols, padded); } + else if (scale) { ggml_cuda_rms_norm_q8(ctx, norm, mul, add, scale, image, cols, padded); } + else { ggml_cuda_rms_norm_q8(ctx, norm, mul, add, scale, image, cols, padded); } +} + +template +static void ggml_cuda_rms_norm_emit_q8_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + ggml_cuda_kernel_launch(rms_norm_emit_q8_f32, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_emit_q8_launch<256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); + } else { + ggml_cuda_rms_norm_emit_q8_launch<1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); + } +} + +void ggml_cuda_op_rms_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_emit_q8(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); } + else if (mul) { ggml_cuda_rms_norm_emit_q8(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); } + else if (scale) { ggml_cuda_rms_norm_emit_q8(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); } + else { ggml_cuda_rms_norm_emit_q8(ctx, norm, mul, add, scale, f16, bf16, image, cols, padded); } +} + +void ggml_cuda_op_hc_post_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, void * f16, void * bf16, void * image, int64_t cols, int64_t padded) { const ggml_tensor * x = post->src[0]; const ggml_tensor * residual = post->src[1]; const ggml_tensor * weights = post->src[2]; @@ -878,10 +1739,163 @@ void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * po const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; - auto kernel = block_size == 256 - ? (comb ? (mul ? hc_post_norm_kernel<256, true, true>(layout) : hc_post_norm_kernel<256, true, false>(layout)) : (mul ? hc_post_norm_kernel<256, false, true>(layout) : hc_post_norm_kernel<256, false, false>(layout))) - : (comb ? (mul ? hc_post_norm_kernel<1024, true, true>(layout) : hc_post_norm_kernel<1024, true, false>(layout)) : (mul ? hc_post_norm_kernel<1024, false, true>(layout) : hc_post_norm_kernel<1024, false, false>(layout))); - ggml_cuda_kernel_launch(kernel, launch, a); + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_post_norm_emit_q8_kernel<256, true, true>(layout) : hc_post_norm_emit_q8_kernel<256, true, false>(layout)) : (mul ? hc_post_norm_emit_q8_kernel<256, false, true>(layout) : hc_post_norm_emit_q8_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_post_norm_emit_q8_kernel<1024, true, true>(layout) : hc_post_norm_emit_q8_kernel<1024, true, false>(layout)) : (mul ? hc_post_norm_emit_q8_kernel<1024, false, true>(layout) : hc_post_norm_emit_q8_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded); +} + +template +static void ggml_cuda_rms_norm_mmq_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + ggml_cuda_kernel_launch(rms_norm_mmq_f32, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), image, cols, padded, rows, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows); + } else { + ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows); + } +} + +template +static void ggml_cuda_rms_norm_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else if (mul) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else if (scale) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } +} + +template +static void ggml_cuda_rms_norm_mxfp4_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + ggml_cuda_kernel_launch(rms_norm_mxfp4_f32< Block, Multiply, Add, Scale>, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), image, cols, padded, rows, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_mxfp4(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_mxfp4_launch< 256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows); + } else { + ggml_cuda_rms_norm_mxfp4_launch< 1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows); + } +} + +static void ggml_cuda_rms_norm_mxfp4_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_mxfp4< true, true, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else if (mul) { ggml_cuda_rms_norm_mxfp4< true, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else if (scale) { ggml_cuda_rms_norm_mxfp4< false, false, true>(ctx, norm, mul, add, scale, image, cols, padded, rows); } + else { ggml_cuda_rms_norm_mxfp4< false, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } +} + +void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout) { + switch (layout) { + case MMQ_Q8_1_DS_LAYOUT_D4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_DS4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; + case 3: ggml_cuda_rms_norm_mxfp4_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; + default: GGML_ABORT("unsupported MMQ image layout"); + } +} + +template +static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + GGML_ASSERT(!mul || !scale); + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = (float *) (scale ? scale : mul ? mul : norm)->data; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm->ne[0]; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = ggml_get_op_params_f32(norm, 0); + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + if (scale) { + auto kernel = block_size == 256 + ? (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) + : (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)); + ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, ggml_get_op_params_f32(scale, 0)); + } else { + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))) + : (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, 1.0f); + } +} + +void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout) { + switch (layout) { + case MMQ_Q8_1_DS_LAYOUT_D4: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_DS4: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; + default: GGML_ABORT("unsupported MMQ image layout"); + } } struct hc_injection_data { @@ -912,49 +1926,6 @@ static __device__ __forceinline__ float hc_injection_gate(hc_injection_data g, i return value; } -struct hc_affine_injection_operand { - const float * data; - uint3 cols; - uint3 rows; - int64_t s0; - int64_t s1; -}; - -struct hc_affine_injection_data { - hc_affine_injection_operand input; - hc_affine_injection_operand multiplier; - hc_affine_injection_operand bias; - uint3 streams; - float * product; - float * added; - float * unary; - float * last; - float scale1; - float bias1; - ggml_unary_op op; - int64_t values; -}; - -static __device__ __forceinline__ float hc_affine_injection_read(hc_affine_injection_operand a, uint32_t col, uint32_t row) { - return a.data[int64_t(fastmodulo(col, a.cols))*a.s0 + int64_t(fastmodulo(row, a.rows))*a.s1]; -} - -static __device__ __forceinline__ float hc_injection_gate(hc_affine_injection_data g, int64_t index, bool store) { - const uint32_t col = fastmodulo(uint32_t(index), g.streams); - const uint32_t row = fastdiv(uint32_t(index), g.streams); - float value = __fmul_rn(hc_affine_injection_read(g.input, col, row), hc_affine_injection_read(g.multiplier, col, row)); - if (store) { g.product[index] = value; } - value = __fadd_rn(value, hc_affine_injection_read(g.bias, col, row)); - if (store) { g.added[index] = value; } - value = g.op == GGML_UNARY_OP_SILU ? ggml_cuda_op_silu_single(value) : 1.0f / (1.0f + expf(-value)); - if (store) { g.unary[index] = value; } - if (g.last) { - value = __fmaf_rn(g.scale1, value, g.bias1); - if (store) { g.last[index] = value; } - } - return value; -} - static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { const int64_t scalar = first + j; @@ -963,99 +1934,6 @@ static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, fl __syncthreads(); } -template -static __device__ __forceinline__ void hc_injection_load(Gate g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { - for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { - const int64_t scalar = first + j; - weights[j] = hc_injection_gate(g, scalar, scalar*n_embd >= begin); - } - __syncthreads(); -} - -template -static __device__ __forceinline__ void hc_injection_norm_impl(hc_post_norm_data a, Gate g) { - ggml_cuda_pdl_lc(); - const int row = blockIdx.x; - const int channel = blockIdx.y; - const int sample = blockIdx.z; - const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; - constexpr bool grouped = layout == 1; - constexpr bool flat = layout == 2; - ggml_cuda_pdl_sync(); - extern __shared__ float s_sum[]; - float * weights = s_sum + 32; - const int64_t gate_first = offset/a.n_embd; - float grouped_weight = 0.0f; - if constexpr (grouped) { - if (threadIdx.x % WARP_SIZE == 0) { - grouped_weight = hc_injection_gate(g, gate_first, threadIdx.x == 0); - } - grouped_weight = __shfl_sync(0xffffffff, grouped_weight, 0, WARP_SIZE); - } else { - const int64_t gate_count = (offset + a.ncols - 1)/a.n_embd - gate_first + 1; - hc_injection_load(g, weights, gate_first, gate_count, offset, a.n_embd); - } - if constexpr (do_multiply) { - a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; - } - const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; - const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; - float tmp = 0.0f; - for (int col = threadIdx.x; col < a.ncols; col += block_size) { - const int64_t ir = offset + col; - const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; - const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; - const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); - const float weight = grouped ? grouped_weight : weights[ih + it*a.hc - gate_first]; - float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weight; - if constexpr (has_comb) { - for (int64_t isrc = 0; isrc < a.hc; ++isrc) { - value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; - } - } else { - value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; - } - a.post_dst[ir] = value; - tmp += value * value; - } - tmp = block_reduce(tmp, s_sum); - const float mean = tmp / a.ncols; - const float scale = rsqrtf(mean + a.eps); - for (int col = threadIdx.x; col < a.ncols; col += block_size) { - if constexpr (do_multiply) { - a.dst[offset + col] = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; - } else { - a.dst[offset + col] = scale * a.post_dst[offset + col]; - } - } -} - -template -static __device__ __forceinline__ void hc_injection_post_impl(hc_post_norm_data a, Gate g) { - ggml_cuda_pdl_lc(); - const int64_t begin = int64_t(blockIdx.x)*blockDim.x; - const int64_t end = begin + int64_t(blockDim.x) < g.values ? begin + int64_t(blockDim.x) : g.values; - const int64_t ir = begin + threadIdx.x; - ggml_cuda_pdl_sync(); - extern __shared__ float weights[]; - const int64_t first = begin/a.n_embd; - const int64_t count = (end - 1)/a.n_embd - first + 1; - hc_injection_load(g, weights, first, count, begin, a.n_embd); - if (ir >= g.values) { return; } - const int64_t i0 = ir % a.n_embd; - const int64_t ih = (ir/a.n_embd) % a.hc; - const int64_t it = ir/(a.n_embd*a.hc); - float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weights[ir/a.n_embd - first]; - if constexpr (has_comb) { - for (int64_t isrc = 0; isrc < a.hc; ++isrc) { - value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; - } - } else { - value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; - } - a.post_dst[ir] = value; -} - template static __global__ __launch_bounds__(block_size) void hc_injection_norm_f32(hc_post_norm_data a, hc_injection_data g) { ggml_cuda_pdl_lc(); @@ -1140,16 +2018,6 @@ static __global__ void hc_injection_post_f32(hc_post_norm_data a, hc_injection_d a.post_dst[ir] = value; } -template -static __global__ __launch_bounds__(block_size) void hc_affine_injection_norm_f32(hc_post_norm_data a, hc_affine_injection_data g) { - hc_injection_norm_impl(a, g); -} - -template -static __global__ void hc_affine_injection_post_f32(hc_post_norm_data a, hc_affine_injection_data g) { - hc_injection_post_impl(a, g); -} - template static auto hc_injection_norm_kernel(int layout) { if (layout == 1) { return hc_injection_norm_f32; } @@ -1215,75 +2083,3 @@ void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * fi : (comb ? (mul ? hc_injection_norm_kernel<1024, true, true>(layout) : hc_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_injection_norm_kernel<1024, false, true>(layout) : hc_injection_norm_kernel<1024, false, false>(layout))); ggml_cuda_kernel_launch(kernel, launch, a, g); } - -static hc_affine_injection_operand ggml_cuda_hc_affine_injection_operand(const ggml_tensor * tensor) { - return {(const float *) tensor->data, init_fastdiv_values(tensor->ne[0]), init_fastdiv_values(tensor->ne[1]), int64_t(tensor->nb[0]/sizeof(float)), int64_t(tensor->nb[1]/sizeof(float))}; -} - -template -static auto hc_affine_injection_norm_kernel(int layout) { - if (layout == 1) { return hc_affine_injection_norm_f32; } - if (layout == 2) { return hc_affine_injection_norm_f32; } - return hc_affine_injection_norm_f32; -} - -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { - const ggml_tensor * x = post->src[0]; - const ggml_tensor * residual = post->src[1]; - const ggml_tensor * weights = post->src[2]; - const ggml_tensor * comb = post->src[3]; - const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; - hc_post_norm_data a{}; - a.x = (const float *) x->data; - a.residual = (const float *) residual->data; - a.post = (const float *) weights->data; - a.comb = comb ? (const float *) comb->data : nullptr; - a.mul = gamma ? (const float *) gamma->data : nullptr; - a.post_dst = (float *) post->data; - a.dst = norm ? (float *) (mul ? mul : norm)->data : nullptr; - a.n_embd = x->ne[0]; - a.hc = residual->ne[1]; - a.ncols = norm ? norm->ne[0] : 256; - a.embd_ne = init_fastdiv_values(a.n_embd); - a.eps = norm ? ggml_get_op_params_f32(norm, 0) : 0.0f; - for (int d = 0; d < 3; ++d) { - if (d < 2) { - a.sx[d] = x->nb[d]/sizeof(float); - a.sp[d] = weights->nb[d]/sizeof(float); - } - a.sr[d] = residual->nb[d]/sizeof(float); - a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; - } - for (int d = 0; d < 4; ++d) { - a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; - a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); - } - hc_affine_injection_data g{}; - g.input = ggml_cuda_hc_affine_injection_operand(first->src[0]); - g.multiplier = ggml_cuda_hc_affine_injection_operand(first->src[1]); - g.bias = ggml_cuda_hc_affine_injection_operand(added->src[1]); - g.streams = init_fastdiv_values(residual->ne[1]); - g.product = (float *) first->data; - g.added = (float *) added->data; - g.unary = (float *) unary->data; - g.last = last != unary ? (float *) last->data : nullptr; - g.scale1 = g.last ? ggml_get_op_params_f32(last, 0) : 1.0f; - g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; - g.op = ggml_get_unary_op(unary); - g.values = ggml_nelements(post); - if (!norm) { - const size_t shared = ((256 + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); - const ggml_cuda_kernel_launch_params launch(dim3((g.values + 255)/256), dim3(256), shared, ctx.stream()); - ggml_cuda_kernel_launch(comb ? hc_affine_injection_post_f32 : hc_affine_injection_post_f32, launch, a, g); - return; - } - const int block_size = a.ncols < 1024 ? 256 : 1024; - const size_t shared = (32 + (a.ncols + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); - const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); - const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); - const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; - auto kernel = block_size == 256 - ? (comb ? (mul ? hc_affine_injection_norm_kernel<256, true, true>(layout) : hc_affine_injection_norm_kernel<256, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<256, false, true>(layout) : hc_affine_injection_norm_kernel<256, false, false>(layout))) - : (comb ? (mul ? hc_affine_injection_norm_kernel<1024, true, true>(layout) : hc_affine_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<1024, false, true>(layout) : hc_affine_injection_norm_kernel<1024, false, false>(layout))); - ggml_cuda_kernel_launch(kernel, launch, a, g); -} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a2bf28bfa16..54522e1f811 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -19,10 +19,31 @@ void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * d void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_rms_norm_emit(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16); + bool ggml_cuda_should_fuse_hc_post_norm(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * mul); void ggml_cuda_op_hc_post_norm(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); -void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); +void ggml_cuda_op_hc_post_norm_emit(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, void * f16, void * bf16); + +void ggml_cuda_op_rms_norm_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded); + +void ggml_cuda_op_rms_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded); + +void ggml_cuda_op_hc_post_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, + ggml_tensor * mul, void * f16, void * bf16, void * image, int64_t cols, int64_t padded); -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); +void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout); + +bool ggml_cuda_should_fuse_hc_post_norm_scale(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * scale); + +void ggml_cuda_op_hc_post_norm_scale(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, + ggml_tensor * scale, void * f16, void * bf16); + +void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout); + +void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); diff --git a/ggml/src/ggml-cuda/reuse.cuh b/ggml/src/ggml-cuda/reuse.cuh new file mode 100644 index 00000000000..5f3232ff1b9 --- /dev/null +++ b/ggml/src/ggml-cuda/reuse.cuh @@ -0,0 +1,194 @@ +#pragma once + +#include "common.cuh" + +#include +#include + +struct ggml_cuda_reuse_pool : ggml_cuda_pool { + char * base; + size_t capacity; + size_t used = 0; + + ggml_cuda_reuse_pool(char * base, size_t capacity) : base(base), capacity(capacity) {} + + void * alloc(size_t size, size_t * actual_size) override { + GGML_ASSERT(size <= SIZE_MAX - 255); + size = GGML_PAD(size, 256); + GGML_ASSERT(size <= capacity - used); + void * ptr = base + used; + used += size; + *actual_size = size; + return ptr; + } + + void free(void * ptr, size_t size) override { + GGML_ASSERT(size <= used); + used -= size; + GGML_ASSERT(ptr == base + used); + } +}; + +struct ggml_cuda_reuse_group { + int node, prepare, last, next; + bool after; + size_t size, offset = 0; +}; + +struct ggml_cuda_reuse_plan { + std::vector groups; + std::vector nodes; + std::vector starts; + std::unordered_map indices; + size_t size = 0; + + ggml_cuda_reuse_plan(const ggml_cgraph * graph, const ggml_cuda_stream_context & streams, + const std::vector & keys, const std::vector & sizes, + bool (*overwritten)(const ggml_tensor *, const ggml_tensor *)) { + if (keys.empty()) { return; } + const int n = graph->n_nodes; + std::unordered_map> consumers; + for (int i = 0; i < n; ++i) { + if (keys[i] >= 0) { consumers[graph->nodes[i]->src[1]].push_back(i); } + } + if (std::none_of(consumers.begin(), consumers.end(), [](const auto & entry) { return entry.second.size() > 1; })) { return; } + nodes.assign(n, -1); + starts.assign(n, -1); + const bool concurrent = !streams.concurrent_events.empty(); + std::vector forks(concurrent ? n : 0, -1), joins(concurrent ? n : 0, -1), lanes(concurrent ? n : 0, 0); + if (!streams.concurrent_events.empty()) { + for (int i = 0; i < n; ++i) { indices[graph->nodes[i]] = i; } + } + for (const auto & entry : streams.concurrent_events) { + const auto first = indices.find(entry.first), last = indices.find(entry.second.join_node); + if (first == indices.end() || last == indices.end() || first->second >= last->second) { + nodes.clear(); starts.clear(); indices.clear(); return; + } + const int fork = first->second, join = last->second; + for (const auto & mapping : entry.second.stream_mapping) { + const auto it = indices.find(mapping.first); + if (it == indices.end()) { nodes.clear(); starts.clear(); indices.clear(); return; } + const int i = it->second; + if (i > fork && i < join) { + forks[i] = fork; + joins[i] = join; + lanes[i] = mapping.second; + } + } + } + const auto ready = [&](const ggml_tensor * tensor, int fork) { + while (tensor->view_src) { tensor = tensor->view_src; } + const auto it = indices.find(tensor); + return it == indices.end() ? tensor->op == GGML_OP_NONE : it->second <= fork; + }; + std::vector members; + groups.reserve(consumers.size()); + for (int i = 0; i < n; ++i) { + if (keys[i] < 0 || nodes[i] >= 0) { continue; } + const ggml_tensor * node = graph->nodes[i]; + const ggml_tensor * input = node->src[1]; + const ggml_tensor * ids = node->op == GGML_OP_MUL_MAT_ID ? node->src[2] : nullptr; + const auto changed = [&](int j) { + return overwritten(graph->nodes[j], input) || (ids && overwritten(graph->nodes[j], ids)); + }; + int prepare = i; + bool after = false; + if (concurrent && forks[i] >= 0) { + const bool early = ready(input, forks[i]) && (!ids || ready(ids, forks[i])); + bool immutable = true; + for (int j = forks[i] + 1; j < joins[i]; ++j) { + if (early || j >= i || lanes[j] != lanes[i]) { immutable &= !changed(j); } + } + if (!immutable) { continue; } + if (early) { prepare = forks[i]; after = true; } + } + members.clear(); + size_t size = 0; + int last = i; + for (int j = prepare + (after ? 1 : 0); j <= consumers.at(input).back(); ++j) { + if (changed(j)) { break; } + const ggml_tensor * other = graph->nodes[j]; + if (j < i || nodes[j] >= 0 || keys[j] != keys[i] || other->src[1] != input) { continue; } + if ((other->op == GGML_OP_MUL_MAT_ID ? other->src[2] : nullptr) != ids || + (ids && other->src[0]->ne[2] != node->src[0]->ne[2])) { continue; } + if (concurrent && forks[i] >= 0 && !after && j < joins[i] && lanes[j] != lanes[i]) { continue; } + if (concurrent && forks[j] >= 0) { + bool immutable = true; + for (int k = forks[j] + 1; k < joins[j]; ++k) { + if (k >= prepare + (after ? 1 : 0) || lanes[k] != lanes[i]) { immutable &= !changed(k); } + } + if (!immutable) { continue; } + } + members.push_back(j); + size = std::max(size, sizes[j]); + last = std::max(last, concurrent && joins[j] >= 0 ? joins[j] : j); + } + if (members.size() < 2 || size > SIZE_MAX - 255) { continue; } + const int group = int(groups.size()); + groups.push_back({i, prepare, last, starts[prepare], after, GGML_PAD(size, 256)}); + starts[prepare] = group; + for (const int member : members) { nodes[member] = group; } + } + pack(); + } + + void pack() { + size = 0; + std::fill(starts.begin(), starts.end(), -1); + for (size_t i = 0; i < groups.size(); ++i) { + auto & group = groups[i]; + group.next = starts[group.prepare]; + starts[group.prepare] = int(i); + } + struct slot { size_t size; int last; }; + std::vector slots; + std::vector order; + order.reserve(groups.size()); + for (size_t i = 0; i < groups.size(); ++i) { order.push_back(int(i)); } + std::sort(order.begin(), order.end(), [&](int a, int b) { return groups[a].prepare < groups[b].prepare; }); + for (const int i : order) { + auto & group = groups[i]; + int best = -1; + for (size_t j = 0; j < slots.size(); ++j) { + if (slots[j].last < group.prepare && (best < 0 || + std::max(slots[j].size, group.size) < std::max(slots[best].size, group.size))) { best = int(j); } + } + if (best < 0) { + best = int(slots.size()); + slots.push_back({group.size, group.last}); + } else { + slots[best].size = std::max(slots[best].size, group.size); + slots[best].last = group.last; + } + group.offset = size_t(best); + } + std::vector offsets; + offsets.reserve(slots.size()); + for (const auto & slot : slots) { + offsets.push_back(size); + if (slot.size > SIZE_MAX - size) { groups.clear(); nodes.clear(); starts.clear(); size = 0; return; } + size += slot.size; + } + for (auto & group : groups) { group.offset = offsets[group.offset]; } + } +}; + +template +struct ggml_cuda_reuse_inputs { + struct item { + ggml_cuda_reuse_pool pool; + T input; + + item(char * base, size_t size) : pool(base, size), input(pool) {} + }; + + ggml_cuda_pool_alloc storage; + std::deque groups; + + ggml_cuda_reuse_inputs(ggml_cuda_pool & pool, const ggml_cuda_reuse_plan & plan) : storage(pool) { + if (plan.size) { storage.alloc(plan.size); } + for (const auto & group : plan.groups) { groups.emplace_back(storage.get() + group.offset, group.size); } + } + + T & operator[](int group) { return groups[group].input; } +}; diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index a90f7df786d..28954cb6970 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -835,3 +835,75 @@ void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_n unary_cuda((const float *)src->data, (float *)sqr_node->data, k, stream); } } + +/* fused scale and activation */ + +template +static __global__ void scaled_unary_kernel(const T * x, T * pre_dst, T * unary_dst, T * dst, int k, float scale0, float bias0, float scale1, float bias1) { + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= k) { + return; + } + ggml_cuda_pdl_sync(); + float value = ggml_cuda_cast(x[i]); + if constexpr (before) { + const T pre = ggml_cuda_cast(scale0 * value + bias0); + pre_dst[i] = pre; + value = ggml_cuda_cast(pre); + } + const T unary = ggml_cuda_cast(op(value)); + unary_dst[i] = unary; + value = ggml_cuda_cast(unary); + if constexpr (after) { + value = scale1 * value + bias1; + } + if constexpr (after) { + dst[i] = ggml_cuda_cast(value); + } +} + +template +static void scaled_unary_cuda(const ggml_tensor * first, const ggml_tensor * unary, const ggml_tensor * last, bool before, bool after, cudaStream_t stream) { + const int k = ggml_nelements(last); + const float scale0 = before ? ggml_get_op_params_f32(first, 0) : 1.0f; + const float bias0 = before ? ggml_get_op_params_f32(first, 1) : 0.0f; + const float scale1 = after ? ggml_get_op_params_f32(last, 0) : 1.0f; + const float bias1 = after ? ggml_get_op_params_f32(last, 1) : 0.0f; + const ggml_cuda_kernel_launch_params launch((k + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE, CUDA_NEG_BLOCK_SIZE, 0, stream); + const T * x = (const T *) first->src[0]->data; + T * pre_dst = before ? (T *) first->data : nullptr; + T * unary_dst = (T *) unary->data; + T * dst = (T *) last->data; + if (before && after) { + ggml_cuda_kernel_launch(scaled_unary_kernel, launch, x, pre_dst, unary_dst, dst, k, scale0, bias0, scale1, bias1); + } else if (before) { + ggml_cuda_kernel_launch(scaled_unary_kernel, launch, x, pre_dst, unary_dst, dst, k, scale0, bias0, scale1, bias1); + } else { + ggml_cuda_kernel_launch(scaled_unary_kernel, launch, x, pre_dst, unary_dst, dst, k, scale0, bias0, scale1, bias1); + } +} + +template +static void ggml_cuda_op_scaled_unary_impl(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last) { + const bool before = first->op == GGML_OP_SCALE; + const bool after = last->op == GGML_OP_SCALE; + if (first->type == GGML_TYPE_BF16) { + scaled_unary_cuda(first, unary, last, before, after, ctx.stream()); + } else { + scaled_unary_cuda(first, unary, last, before, after, ctx.stream()); + } +} + +void ggml_cuda_op_scaled_unary(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last) { + switch (ggml_get_unary_op(unary)) { + case GGML_UNARY_OP_SILU: + ggml_cuda_op_scaled_unary_impl(ctx, first, unary, last); + break; + case GGML_UNARY_OP_SIGMOID: + ggml_cuda_op_scaled_unary_impl(ctx, first, unary, last); + break; + default: + GGML_ABORT("Unsupported scaled unary op"); + } +} diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 8439d85afeb..dafb4c46f98 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -123,3 +123,5 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl return ggml_cuda_op_silu_single(gate) * up; } + +void ggml_cuda_op_scaled_unary(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last); From f04864e2ad6a080ed1cdd5051d6aee0fe7d560ac Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 02:31:16 -0400 Subject: [PATCH 04/27] cuda : preserve prepared images during affine HC fusion Reuse the allocated affine match to find the planned HC_POST and retain prepared emission priority. Reject expanded normalization cuts whose terminal is not written by this kernel. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 185 ++++++++++++++++++++++++++- ggml/src/ggml-cuda/norm.cu | 218 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/norm.cuh | 2 + 3 files changed, 403 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 625f1595af3..d7357c6b4e6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4919,8 +4919,147 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu return match; } +struct ggml_cuda_hc_affine_injection_match { + ggml_cuda_affine_unary_match affine; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + ggml_cuda_hc_post_norm_match hc; + int count = 0; +}; + +static bool ggml_cuda_hc_affine_tensor_ok(const ggml_tensor * tensor) { + if (!tensor || tensor->type != GGML_TYPE_F32 || (tensor->view_src && tensor->view_src->type != GGML_TYPE_F32)) { return false; } + size_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] <= 0 || tensor->ne[d] > INT_MAX || tensor->nb[d] % sizeof(float) || elements > INT64_MAX/sizeof(float)/size_t(tensor->ne[d])) { return false; } + elements *= tensor->ne[d]; + } + size_t bytes; + return ggml_cuda_affine_unary_bytes(tensor, bytes) && bytes <= INT64_MAX; +} + +static bool ggml_cuda_hc_affine_prior_source(ggml_cgraph * graph, int start, const ggml_tensor * tensor) { + if (!tensor) { return false; } + if (tensor->view_src) { return ggml_cuda_affine_unary_prior_view(graph, start, tensor); } + if (tensor->op == GGML_OP_NONE) { return true; } + for (int i = 0; i < start; ++i) { + if (graph->nodes[i] == tensor) { return true; } + } + return false; +} + +static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + auto affine = ggml_cuda_match_affine_unary(ctx, graph, i, allocated); + if (!affine.count || affine.mul->type != GGML_TYPE_F32 || affine.add->type != GGML_TYPE_F32 || affine.unary->type != GGML_TYPE_F32) { return {}; } + ggml_tensor * last = affine.post ? affine.post : affine.unary; + const ggml_tensor * gate_inputs[] = {affine.mul->src[0], affine.mul->src[1], affine.add->src[1]}; + for (const ggml_tensor * tensor : gate_inputs) { + if (!ggml_cuda_hc_affine_tensor_ok(tensor) || tensor->ne[2] != 1 || tensor->ne[3] != 1 || affine.mul->ne[0] % tensor->ne[0] || affine.mul->ne[1] % tensor->ne[1]) { return {}; } + } + int post_i = i + affine.count; + auto skip_views = [&]() { + while (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_VIEW) { + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[post_i])) { return false; } + ++post_i; + } + return true; + }; + if (!skip_views()) { return {}; } + ggml_tensor * comb = nullptr; + if (post_i < graph->n_nodes && graph->nodes[post_i]->op == GGML_OP_DSV4_HC_COMB) { + comb = graph->nodes[post_i++]; + // The existing COMB operator supports four streams. + if (comb->type != GGML_TYPE_F32 || comb->ne[0] != 4 || comb->ne[1] != 4 || comb->ne[3] != 1) { return {}; } + const int64_t mix_dim = (comb->ne[0] + 2)*comb->ne[0]; + const ggml_tensor * mixes = comb->src[0]; + const ggml_tensor * scale = comb->src[1]; + const ggml_tensor * base = comb->src[2]; + if (!mixes || !scale || !base || mixes->ne[0] != mix_dim || mixes->ne[1] != comb->ne[2] || mixes->ne[2] != 1 || mixes->ne[3] != 1 || + scale->ne[0] < 3 || scale->ne[1] != 1 || scale->ne[2] != 1 || scale->ne[3] != 1 || + base->ne[0] != mix_dim || base->ne[1] != 1 || base->ne[2] != 1 || base->ne[3] != 1 || comb->op_params[1] <= 0) { return {}; } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (comb->src[j] && (!ggml_cuda_hc_affine_tensor_ok(comb->src[j]) || !ggml_cuda_hc_affine_prior_source(graph, i, comb->src[j]))) { return {}; } + } + if (!skip_views()) { return {}; } + } + if (post_i >= graph->n_nodes) { return {}; } + ggml_tensor * post = graph->nodes[post_i]; + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * post_comb = post->src[3]; + if (post->op != GGML_OP_DSV4_HC_POST || post->src[2] != last || !x || !residual || + x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || + !ggml_is_contiguous(post) || !ggml_are_same_shape(post, residual) || x->ne[0] <= 0 || x->ne[1] <= 0 || residual->ne[1] <= 0 || + x->ne[2] != 1 || x->ne[3] != 1 || residual->ne[0] != x->ne[0] || residual->ne[2] != x->ne[1] || residual->ne[3] != 1 || + last->ne[0] != residual->ne[1] || last->ne[1] != x->ne[1] || last->ne[2] != 1 || last->ne[3] != 1 || + (comb && post_comb != comb) || + (post_comb && (post_comb->type != GGML_TYPE_F32 || post_comb->ne[0] != residual->ne[1] || post_comb->ne[1] != residual->ne[1] || post_comb->ne[2] != x->ne[1] || post_comb->ne[3] != 1))) { return {}; } + for (const ggml_tensor * tensor : {x, residual, (const ggml_tensor *) post, post_comb}) { + if (tensor && !ggml_cuda_hc_affine_tensor_ok(tensor)) { return {}; } + } + size_t bytes; + if (!ggml_cuda_hc_injection_bytes(post, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } + auto hc = ggml_cuda_match_hc_post_norm(graph, post_i); + if (hc.count && graph->nodes[post_i + hc.count - 1] != (hc.mul ? hc.mul : hc.norm)) { return {}; } + if (hc.count && hc.norm->ne[0] != x->ne[0]) { + if (post_comb && bytes/sizeof(float) <= 256) { return {}; } + if (hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] > WARP_SIZE/2 && hc.norm->ne[0] < WARP_SIZE) { return {}; } + const int block_size = hc.norm->ne[0] < 1024 ? 256 : 1024; + const int64_t nweights = (hc.norm->ne[0] + x->ne[0] - 2)/x->ne[0] + 1; + if ((hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] <= WARP_SIZE/2) || nweights >= block_size) { hc = {}; } + } + const int end = hc.count ? post_i + hc.count : post_i + 1; + const int64_t ncols = hc.count ? hc.norm->ne[0] : 256; + const size_t shared = ((hc.count ? 32 : 0) + (ncols + x->ne[0] - 2)/x->ne[0] + 1)*sizeof(float); + if (shared > ggml_cuda_info().devices[ctx.device].smpb) { return {}; } + std::vector ops; + std::vector outputs; + for (int j = i; j < end; ++j) { + ops.push_back(graph->nodes[j]->op); + if (!hc.mul || graph->nodes[j] != hc.norm) { outputs.push_back(j); } + } + if (!ggml_can_fuse_subgraph(graph, i, end - i, ops.data(), outputs.data(), outputs.size())) { return {}; } + if (allocated) { + if (hc.count && !ggml_cuda_hc_post_norm_memory_ok(graph, post_i, hc, ctx.device)) { return {}; } + std::vector reads = {gate_inputs[0], gate_inputs[1], gate_inputs[2], x, residual, post_comb}; + if (hc.mul) { reads.push_back(hc.mul->src[hc.mul->src[0] == hc.norm ? 1 : 0]); } + if (comb) { + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (comb->src[j]) { reads.push_back(comb->src[j]); } + } + } + std::vector writes = {affine.mul, affine.add, affine.unary}; + if (affine.post) { writes.push_back(affine.post); } + writes.push_back(post); + if (hc.count) { writes.push_back(hc.mul ? hc.mul : hc.norm); } + if (comb) { writes.push_back(comb); } + std::vector rb(reads.size()), re(reads.size()), wb(writes.size()), we(writes.size()); + for (size_t j = 0; j < reads.size(); ++j) { + if (reads[j] && (!ggml_cuda_hc_affine_tensor_ok(reads[j]) || !ggml_cuda_hc_injection_range(reads[j], ctx.device, rb[j], re[j]))) { return {}; } + } + for (size_t j = 0; j < writes.size(); ++j) { + if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_is_contiguous(writes[j]) || !ggml_cuda_hc_affine_tensor_ok(writes[j]) || !ggml_cuda_hc_injection_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } + for (size_t k = 0; k < reads.size(); ++k) { + // COMB is written before the affine kernel reads it. + if (reads[k] && !(writes[j] == comb && reads[k] == comb) && wb[j] < re[k] && rb[k] < we[j]) { return {}; } + } + for (size_t k = 0; k < j; ++k) { + if (wb[j] < we[k] && wb[k] < we[j]) { return {}; } + } + } + for (int j = i; j < end; ++j) { + const ggml_tensor * tensor = graph->nodes[j]; + if (tensor->op != GGML_OP_VIEW && tensor->op != GGML_OP_RESHAPE) { continue; } + uintptr_t begin, stop; + if (!ggml_cuda_hc_injection_range(tensor, ctx.device, begin, stop)) { return {}; } + } + } + return {affine, post, comb, hc, end - i}; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, - const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1) { + const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, + const ggml_cuda_hc_affine_injection_match * affine_match = nullptr) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { @@ -5716,6 +5855,19 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto injection = affine_match ? *affine_match : ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); + if (injection.count && (prepared_hc_post < 0 || injection.post != cgraph->nodes[prepared_hc_post] || (!injection.hc.norm && injection.post->src[3]))) { + if (injection.comb) { + const bool ok = ggml_cuda_compute_forward(*cuda_ctx, injection.comb); + GGML_ASSERT(ok); + } + const auto & affine = injection.affine; + ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, injection.post, injection.hc.norm, injection.hc.mul); + return injection.count - 1; + } + } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); if (affine.count) { @@ -6127,7 +6279,20 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } } - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post); + ggml_cuda_hc_affine_injection_match affine_injection; + const ggml_cuda_hc_affine_injection_match * affine_match = nullptr; + if (node->op == GGML_OP_MUL && cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty() && + (!norm_emits.empty() || !q8_emits.empty() || !mmq_emits.empty())) { + affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); + affine_match = &affine_injection; + if (affine_injection.count) { + const int post_i = i + affine_injection.count - (affine_injection.hc.count ? affine_injection.hc.count : 1); + if ((!norm_emits.empty() && norm_emits[post_i].post) || (!q8_emits.empty() && q8_emits[post_i].post) || (!mmq_emits.empty() && mmq_emits[post_i].post)) { + prepared_hc_post = post_i; + } + } + } + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match); if (nodes_to_skip != 0) { #ifdef GGML_CUDA_DEBUG @@ -6343,6 +6508,22 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i, false); + if (affine_injection.count) { + ggml_tensor * terminal = cgraph->nodes[i + affine_injection.count - 1]; + ggml_tensor * unused_norm = affine_injection.hc.mul ? affine_injection.hc.norm : nullptr; + for (int j = i; j < i + affine_injection.count - 1; ++j) { + if (cgraph->nodes[j] != unused_norm) { params->add_alloc_dep(params->user_data, cgraph->nodes[j], terminal); } + for (int k = 0; k < GGML_MAX_SRC; ++k) { + if (cgraph->nodes[j]->src[k]) { params->add_alloc_dep(params->user_data, cgraph->nodes[j]->src[k], terminal); } + } + } + for (ggml_tensor * src : terminal->src) { + if (src && src != unused_norm) { params->add_alloc_dep(params->user_data, src, terminal); } + } + i += affine_injection.count - 1; + continue; + } const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); if (affine.count) { ggml_tensor * last = cgraph->nodes[i + affine.count - 1]; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 355fc21e493..cdfdb0769b4 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1926,6 +1926,49 @@ static __device__ __forceinline__ float hc_injection_gate(hc_injection_data g, i return value; } +struct hc_affine_injection_operand { + const float * data; + uint3 cols; + uint3 rows; + int64_t s0; + int64_t s1; +}; + +struct hc_affine_injection_data { + hc_affine_injection_operand input; + hc_affine_injection_operand multiplier; + hc_affine_injection_operand bias; + uint3 streams; + float * product; + float * added; + float * unary; + float * last; + float scale1; + float bias1; + ggml_unary_op op; + int64_t values; +}; + +static __device__ __forceinline__ float hc_affine_injection_read(hc_affine_injection_operand a, uint32_t col, uint32_t row) { + return a.data[int64_t(fastmodulo(col, a.cols))*a.s0 + int64_t(fastmodulo(row, a.rows))*a.s1]; +} + +static __device__ __forceinline__ float hc_injection_gate(hc_affine_injection_data g, int64_t index, bool store) { + const uint32_t col = fastmodulo(uint32_t(index), g.streams); + const uint32_t row = fastdiv(uint32_t(index), g.streams); + float value = __fmul_rn(hc_affine_injection_read(g.input, col, row), hc_affine_injection_read(g.multiplier, col, row)); + if (store) { g.product[index] = value; } + value = __fadd_rn(value, hc_affine_injection_read(g.bias, col, row)); + if (store) { g.added[index] = value; } + value = g.op == GGML_UNARY_OP_SILU ? ggml_cuda_op_silu_single(value) : 1.0f / (1.0f + expf(-value)); + if (store) { g.unary[index] = value; } + if (g.last) { + value = __fmaf_rn(g.scale1, value, g.bias1); + if (store) { g.last[index] = value; } + } + return value; +} + static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { const int64_t scalar = first + j; @@ -1934,6 +1977,99 @@ static __device__ __forceinline__ void hc_injection_load(hc_injection_data g, fl __syncthreads(); } +template +static __device__ __forceinline__ void hc_injection_load(Gate g, float * weights, int64_t first, int64_t count, int64_t begin, int64_t n_embd) { + for (int64_t j = threadIdx.x; j < count; j += blockDim.x) { + const int64_t scalar = first + j; + weights[j] = hc_injection_gate(g, scalar, scalar*n_embd >= begin); + } + __syncthreads(); +} + +template +static __device__ __forceinline__ void hc_injection_norm_impl(hc_post_norm_data a, Gate g) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + ggml_cuda_pdl_sync(); + extern __shared__ float s_sum[]; + float * weights = s_sum + 32; + const int64_t gate_first = offset/a.n_embd; + float grouped_weight = 0.0f; + if constexpr (grouped) { + if (threadIdx.x % WARP_SIZE == 0) { + grouped_weight = hc_injection_gate(g, gate_first, threadIdx.x == 0); + } + grouped_weight = __shfl_sync(0xffffffff, grouped_weight, 0, WARP_SIZE); + } else { + const int64_t gate_count = (offset + a.ncols - 1)/a.n_embd - gate_first + 1; + hc_injection_load(g, weights, gate_first, gate_count, offset, a.n_embd); + } + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; + } + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + const float weight = grouped ? grouped_weight : weights[ih + it*a.hc - gate_first]; + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weight; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + a.dst[offset + col] = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + } else { + a.dst[offset + col] = scale * a.post_dst[offset + col]; + } + } +} + +template +static __device__ __forceinline__ void hc_injection_post_impl(hc_post_norm_data a, Gate g) { + ggml_cuda_pdl_lc(); + const int64_t begin = int64_t(blockIdx.x)*blockDim.x; + const int64_t end = begin + int64_t(blockDim.x) < g.values ? begin + int64_t(blockDim.x) : g.values; + const int64_t ir = begin + threadIdx.x; + ggml_cuda_pdl_sync(); + extern __shared__ float weights[]; + const int64_t first = begin/a.n_embd; + const int64_t count = (end - 1)/a.n_embd - first + 1; + hc_injection_load(g, weights, first, count, begin, a.n_embd); + if (ir >= g.values) { return; } + const int64_t i0 = ir % a.n_embd; + const int64_t ih = (ir/a.n_embd) % a.hc; + const int64_t it = ir/(a.n_embd*a.hc); + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weights[ir/a.n_embd - first]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; +} + template static __global__ __launch_bounds__(block_size) void hc_injection_norm_f32(hc_post_norm_data a, hc_injection_data g) { ggml_cuda_pdl_lc(); @@ -2018,6 +2154,16 @@ static __global__ void hc_injection_post_f32(hc_post_norm_data a, hc_injection_d a.post_dst[ir] = value; } +template +static __global__ __launch_bounds__(block_size) void hc_affine_injection_norm_f32(hc_post_norm_data a, hc_affine_injection_data g) { + hc_injection_norm_impl(a, g); +} + +template +static __global__ void hc_affine_injection_post_f32(hc_post_norm_data a, hc_affine_injection_data g) { + hc_injection_post_impl(a, g); +} + template static auto hc_injection_norm_kernel(int layout) { if (layout == 1) { return hc_injection_norm_f32; } @@ -2083,3 +2229,75 @@ void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * fi : (comb ? (mul ? hc_injection_norm_kernel<1024, true, true>(layout) : hc_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_injection_norm_kernel<1024, false, true>(layout) : hc_injection_norm_kernel<1024, false, false>(layout))); ggml_cuda_kernel_launch(kernel, launch, a, g); } + +static hc_affine_injection_operand ggml_cuda_hc_affine_injection_operand(const ggml_tensor * tensor) { + return {(const float *) tensor->data, init_fastdiv_values(tensor->ne[0]), init_fastdiv_values(tensor->ne[1]), int64_t(tensor->nb[0]/sizeof(float)), int64_t(tensor->nb[1]/sizeof(float))}; +} + +template +static auto hc_affine_injection_norm_kernel(int layout) { + if (layout == 1) { return hc_affine_injection_norm_f32; } + if (layout == 2) { return hc_affine_injection_norm_f32; } + return hc_affine_injection_norm_f32; +} + +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { + const ggml_tensor * x = post->src[0]; + const ggml_tensor * residual = post->src[1]; + const ggml_tensor * weights = post->src[2]; + const ggml_tensor * comb = post->src[3]; + const ggml_tensor * gamma = mul ? mul->src[mul->src[0] == norm ? 1 : 0] : nullptr; + hc_post_norm_data a{}; + a.x = (const float *) x->data; + a.residual = (const float *) residual->data; + a.post = (const float *) weights->data; + a.comb = comb ? (const float *) comb->data : nullptr; + a.mul = gamma ? (const float *) gamma->data : nullptr; + a.post_dst = (float *) post->data; + a.dst = norm ? (float *) (mul ? mul : norm)->data : nullptr; + a.n_embd = x->ne[0]; + a.hc = residual->ne[1]; + a.ncols = norm ? norm->ne[0] : 256; + a.embd_ne = init_fastdiv_values(a.n_embd); + a.eps = norm ? ggml_get_op_params_f32(norm, 0) : 0.0f; + for (int d = 0; d < 3; ++d) { + if (d < 2) { + a.sx[d] = x->nb[d]/sizeof(float); + a.sp[d] = weights->nb[d]/sizeof(float); + } + a.sr[d] = residual->nb[d]/sizeof(float); + a.sc[d] = comb ? comb->nb[d]/sizeof(float) : 0; + } + for (int d = 0; d < 4; ++d) { + a.sm[d] = gamma ? gamma->nb[d]/sizeof(float) : 0; + a.mul_ne[d] = init_fastdiv_values(gamma ? gamma->ne[d] : 1); + } + hc_affine_injection_data g{}; + g.input = ggml_cuda_hc_affine_injection_operand(first->src[0]); + g.multiplier = ggml_cuda_hc_affine_injection_operand(first->src[1]); + g.bias = ggml_cuda_hc_affine_injection_operand(added->src[1]); + g.streams = init_fastdiv_values(residual->ne[1]); + g.product = (float *) first->data; + g.added = (float *) added->data; + g.unary = (float *) unary->data; + g.last = last != unary ? (float *) last->data : nullptr; + g.scale1 = g.last ? ggml_get_op_params_f32(last, 0) : 1.0f; + g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; + g.op = ggml_get_unary_op(unary); + g.values = ggml_nelements(post); + if (!norm) { + const size_t shared = ((256 + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const ggml_cuda_kernel_launch_params launch(dim3((g.values + 255)/256), dim3(256), shared, ctx.stream()); + ggml_cuda_kernel_launch(comb ? hc_affine_injection_post_f32 : hc_affine_injection_post_f32, launch, a, g); + return; + } + const int block_size = a.ncols < 1024 ? 256 : 1024; + const size_t shared = (32 + (a.ncols + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + auto kernel = block_size == 256 + ? (comb ? (mul ? hc_affine_injection_norm_kernel<256, true, true>(layout) : hc_affine_injection_norm_kernel<256, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<256, false, true>(layout) : hc_affine_injection_norm_kernel<256, false, false>(layout))) + : (comb ? (mul ? hc_affine_injection_norm_kernel<1024, true, true>(layout) : hc_affine_injection_norm_kernel<1024, true, false>(layout)) : (mul ? hc_affine_injection_norm_kernel<1024, false, true>(layout) : hc_affine_injection_norm_kernel<1024, false, false>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, g); +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 54522e1f811..4a06fdc8287 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -47,3 +47,5 @@ void ggml_cuda_op_hc_post_norm_scale(ggml_backend_cuda_context & ctx, ggml_tenso void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout); void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); + +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); From c2de3b5d10e925674bdea5e99169f70c95dc8122 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 03:47:46 -0400 Subject: [PATCH 05/27] cuda : fold affine HC gates into prepared image writers Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 105 ++++++++++++++++++++-------- ggml/src/ggml-cuda/norm.cu | 119 +++++++++++++++++++++++++++++++- ggml/src/ggml-cuda/norm.cuh | 13 +++- 3 files changed, 205 insertions(+), 32 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index d7357c6b4e6..e0009d3605c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5059,7 +5059,7 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, - const ggml_cuda_hc_affine_injection_match * affine_match = nullptr) { + const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { @@ -5857,6 +5857,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto injection = affine_match ? *affine_match : ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); + if (injection.count && affine_emit) { + *affine_emit = true; + return 0; + } if (injection.count && (prepared_hc_post < 0 || injection.post != cgraph->nodes[prepared_hc_post] || (!injection.hc.norm && injection.post->src[3]))) { if (injection.comb) { const bool ok = ggml_cuda_compute_forward(*cuda_ctx, injection.comb); @@ -6159,9 +6163,51 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } const ggml_tensor * shared_input = q8_group >= 0 ? node->src[1] : nullptr; const char * quantized = q8_group >= 0 ? q8_inputs[q8_group].get() : nullptr; - if (!mmq_emits.empty() && mmq_emits[i].post && (q8_emits.empty() || !q8_emits[i].norm)) { - const auto & emit = mmq_emits[i]; - const auto * typed = !norm_emits.empty() && norm_emits[i].norm ? &norm_emits[i] : nullptr; + int prepared_hc_post = -1; + ggml_cuda_hc_affine_injection_match affine_injection; + const ggml_cuda_hc_affine_injection_match * affine_match = nullptr; + if (node->op == GGML_OP_MUL && cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty() && + (!norm_emits.empty() || !q8_emits.empty() || !mmq_emits.empty())) { + affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); + affine_match = &affine_injection; + if (affine_injection.count) { + const int post_i = i + affine_injection.count - (affine_injection.hc.count ? affine_injection.hc.count : 1); + if ((!norm_emits.empty() && norm_emits[post_i].post) || (!q8_emits.empty() && q8_emits[post_i].post) || (!mmq_emits.empty() && mmq_emits[post_i].post)) { + prepared_hc_post = post_i; + } + } + } + bool affine_emit = false; + if (affine_match && affine_injection.hc.norm && prepared_hc_post >= 0) { + const auto fits = [&](const auto & emit) { + return !emit.norm || (emit.post == affine_injection.post && emit.norm == affine_injection.hc.norm && + emit.mul == affine_injection.hc.mul && !emit.add && !emit.scale && + emit.dst == (emit.mul ? emit.mul : emit.norm) && emit.last == i + affine_injection.count - 1); + }; + if ((norm_emits.empty() || fits(norm_emits[prepared_hc_post])) && + (q8_emits.empty() || fits(q8_emits[prepared_hc_post])) && + (mmq_emits.empty() || fits(mmq_emits[prepared_hc_post]))) { + const int skipped = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match, &affine_emit); + if (skipped) { + i += skipped; + continue; + } + GGML_ASSERT(affine_emit); + } + } + const int emit_i = affine_emit ? prepared_hc_post : i; + const auto launch_affine_emit = [&](const ggml_cuda_hc_affine_emit_data & data) { + if (affine_injection.comb) { + const bool ok = ggml_cuda_compute_forward(*cuda_ctx, affine_injection.comb); + GGML_ASSERT(ok); + } + const auto & affine = affine_injection.affine; + ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, + affine_injection.post, affine_injection.hc.norm, affine_injection.hc.mul, &data); + }; + if (!mmq_emits.empty() && mmq_emits[emit_i].post && (q8_emits.empty() || !q8_emits[emit_i].norm)) { + const auto & emit = mmq_emits[emit_i]; + const auto * typed = !norm_emits.empty() && norm_emits[emit_i].norm ? &norm_emits[emit_i] : nullptr; if (!typed || (typed->post == emit.post && typed->norm == emit.norm && typed->dst == emit.dst && typed->last == emit.last)) { const int g = emit.image; const int reader = mmq_reuse.groups[g].node; @@ -6174,23 +6220,27 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud }; void * f16 = typed ? image(typed->f16) : nullptr; void * bf16 = typed ? image(typed->bf16) : nullptr; - ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, - mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + if (affine_emit) { + launch_affine_emit({f16, bf16, nullptr, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]}); + } else { + ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + } i = emit.last; continue; } } - if (!q8_emits.empty() && q8_emits[i].norm && - ((!norm_emits.empty() && norm_emits[i].norm) || q8_emits[i].post)) { - const auto & emit = q8_emits[i]; + if (!q8_emits.empty() && q8_emits[emit_i].norm && + ((!norm_emits.empty() && norm_emits[emit_i].norm) || q8_emits[emit_i].post)) { + const auto & emit = q8_emits[emit_i]; const int g = emit.image; const ggml_tensor * input = cgraph->nodes[q8_reuse.groups[g].node]->src[1]; q8_inputs[g].alloc(mmvq_sizes[q8_reuse.groups[g].node]); void * f16 = nullptr; void * bf16 = nullptr; - if (!norm_emits.empty() && norm_emits[i].norm) { - const auto & typed = norm_emits[i]; + if (!norm_emits.empty() && norm_emits[emit_i].norm) { + const auto & typed = norm_emits[emit_i]; GGML_ASSERT(typed.norm == emit.norm && typed.dst == emit.dst && typed.last == emit.last); const auto image = [&](int group) -> void * { if (group < 0) { return nullptr; } @@ -6206,8 +6256,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_op_rms_norm_emit_q8(*cuda_ctx, emit.norm, nullptr, nullptr, emit.scale, f16, bf16, q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); } else { - ggml_cuda_op_hc_post_norm_emit_q8(*cuda_ctx, emit.post, emit.norm, emit.mul, f16, bf16, - q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + if (affine_emit) { + launch_affine_emit({f16, bf16, q8_inputs[g].get(), nullptr, input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], 0}); + } else { + ggml_cuda_op_hc_post_norm_emit_q8(*cuda_ctx, emit.post, emit.norm, emit.mul, f16, bf16, + q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); + } } } else { ggml_cuda_op_rms_norm_emit_q8(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, f16, bf16, @@ -6217,8 +6271,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } - if (!norm_emits.empty() && norm_emits[i].norm) { - const auto & emit = norm_emits[i]; + if (!norm_emits.empty() && norm_emits[emit_i].norm) { + const auto & emit = norm_emits[emit_i]; const auto image = [&](int g) -> void * { if (g < 0) { return nullptr; } shared_inputs[g].alloc(input_sizes[reuse.groups[g].node]); @@ -6228,7 +6282,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (emit.scale) { ggml_cuda_op_hc_post_norm_scale(*cuda_ctx, emit.post, emit.norm, emit.scale, image(emit.f16), image(emit.bf16)); } else { - ggml_cuda_op_hc_post_norm_emit(*cuda_ctx, emit.post, emit.norm, emit.mul, image(emit.f16), image(emit.bf16)); + if (affine_emit) { + launch_affine_emit({image(emit.f16), image(emit.bf16)}); + } else { + ggml_cuda_op_hc_post_norm_emit(*cuda_ctx, emit.post, emit.norm, emit.mul, image(emit.f16), image(emit.bf16)); + } } } else { ggml_cuda_op_rms_norm_emit(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, image(emit.f16), image(emit.bf16)); @@ -6270,7 +6328,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } if (ready) { prepare_mmq_group(mmq_group); } } - int prepared_hc_post = -1; + if (node->op == GGML_OP_SCALE || node->op == GGML_OP_UNARY) { for (int j = i + 1; j < cgraph->n_nodes && j - i <= 3; ++j) { if ((!norm_emits.empty() && norm_emits[j].post) || (!q8_emits.empty() && q8_emits[j].post) || (!mmq_emits.empty() && mmq_emits[j].post)) { @@ -6279,19 +6337,6 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } } - ggml_cuda_hc_affine_injection_match affine_injection; - const ggml_cuda_hc_affine_injection_match * affine_match = nullptr; - if (node->op == GGML_OP_MUL && cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty() && - (!norm_emits.empty() || !q8_emits.empty() || !mmq_emits.empty())) { - affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i); - affine_match = &affine_injection; - if (affine_injection.count) { - const int post_i = i + affine_injection.count - (affine_injection.hc.count ? affine_injection.hc.count : 1); - if ((!norm_emits.empty() && norm_emits[post_i].post) || (!q8_emits.empty() && q8_emits[post_i].post) || (!mmq_emits.empty() && mmq_emits[post_i].post)) { - prepared_hc_post = post_i; - } - } - } int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match); if (nodes_to_skip != 0) { diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index cdfdb0769b4..985d34ba4ff 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -2234,6 +2234,105 @@ static hc_affine_injection_operand ggml_cuda_hc_affine_injection_operand(const g return {(const float *) tensor->data, init_fastdiv_values(tensor->ne[0]), init_fastdiv_values(tensor->ne[1]), int64_t(tensor->nb[0]/sizeof(float)), int64_t(tensor->nb[1]/sizeof(float))}; } +template +static __device__ __forceinline__ void hc_injection_norm_emit_impl(hc_post_norm_data a, Gate g, const Write write) { + ggml_cuda_pdl_lc(); + const int row = blockIdx.x; + const int channel = blockIdx.y; + const int sample = blockIdx.z; + const int64_t offset = ((int64_t(sample)*gridDim.y + channel)*gridDim.x + row)*a.ncols; + constexpr bool grouped = layout == 1; + constexpr bool flat = layout == 2; + ggml_cuda_pdl_sync(); + extern __shared__ float s_sum[]; + float * weights = s_sum + 32; + const int64_t gate_first = offset/a.n_embd; + float grouped_weight = 0.0f; + if constexpr (grouped) { + if (threadIdx.x % WARP_SIZE == 0) { + grouped_weight = hc_injection_gate(g, gate_first, threadIdx.x == 0); + } + grouped_weight = __shfl_sync(0xffffffff, grouped_weight, 0, WARP_SIZE); + } else { + const int64_t gate_count = (offset + a.ncols - 1)/a.n_embd - gate_first + 1; + hc_injection_load(g, weights, gate_first, gate_count, offset, a.n_embd); + } + if constexpr (do_multiply) { + a.mul += fastmodulo(sample, a.mul_ne[3])*a.sm[3] + fastmodulo(channel, a.mul_ne[2])*a.sm[2] + fastmodulo(row, a.mul_ne[1])*a.sm[1]; + } + const int64_t stream = grouped ? (offset / a.n_embd) % a.hc : 0; + const int64_t token = grouped || flat ? offset / (a.n_embd*a.hc) : 0; + float tmp = 0.0f; + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + const int64_t ir = offset + col; + const int64_t i0 = grouped ? col : flat ? fastmodulo(col, a.embd_ne) : ir % a.n_embd; + const int64_t ih = grouped ? stream : flat ? fastdiv(col, a.embd_ne) : (ir / a.n_embd) % a.hc; + const int64_t it = grouped || flat ? token : ir / (a.n_embd*a.hc); + const float weight = grouped ? grouped_weight : weights[ih + it*a.hc - gate_first]; + float value = a.x[i0*a.sx[0] + it*a.sx[1]] * weight; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < a.hc; ++isrc) { + value += a.residual[i0*a.sr[0] + isrc*a.sr[1] + it*a.sr[2]] * a.comb[ih*a.sc[0] + isrc*a.sc[1] + it*a.sc[2]]; + } + } else { + value += a.residual[i0*a.sr[0] + ih*a.sr[1] + it*a.sr[2]]; + } + a.post_dst[ir] = value; + tmp += value * value; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / a.ncols; + const float scale = rsqrtf(mean + a.eps); + if constexpr (Write::width == 2) { + for (int64_t col = 2*threadIdx.x; col < a.ncols; col += 2*block_size) { + float x, y; + if constexpr (do_multiply) { + x = scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]; + y = scale * a.post_dst[offset + col + 1] * a.mul[fastmodulo(col + 1, a.mul_ne[0])]; + } else { + x = scale * a.post_dst[offset + col]; + y = scale * a.post_dst[offset + col + 1]; + } + write(a.dst + offset, a.dst, int(col), make_float2(x, y)); + } + } else { + for (int col = threadIdx.x; col < a.ncols; col += block_size) { + if constexpr (do_multiply) { + write(a.dst + offset, a.dst, col, scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])]); + } else { + write(a.dst + offset, a.dst, col, scale * a.post_dst[offset + col]); + } + } + } +} + + +template +static __global__ __launch_bounds__(Block) void hc_affine_injection_emit_f32(hc_post_norm_data a, hc_affine_injection_data g, Write write) { + hc_injection_norm_emit_impl(a, g, write); +} + +template +static auto hc_affine_injection_emit_kernel(int layout) { + if (layout == 1) { return hc_affine_injection_emit_f32; } + if (layout == 2) { return hc_affine_injection_emit_f32; } + return hc_affine_injection_emit_f32; +} + +template +static void ggml_cuda_hc_affine_emit_launch(ggml_backend_cuda_context & ctx, hc_post_norm_data a, hc_affine_injection_data g, + ggml_tensor * norm, bool comb, bool multiply, const Write write) { + const int block_size = a.ncols < 1024 ? 256 : 1024; + const size_t shared = (32 + (a.ncols + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); + const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); + const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); + const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + auto kernel = block_size == 256 + ? (comb ? (multiply ? hc_affine_injection_emit_kernel<256, true, true, Write>(layout) : hc_affine_injection_emit_kernel<256, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<256, false, true, Write>(layout) : hc_affine_injection_emit_kernel<256, false, false, Write>(layout))) + : (comb ? (multiply ? hc_affine_injection_emit_kernel<1024, true, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<1024, false, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, false, false, Write>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, g, write); +} + template static auto hc_affine_injection_norm_kernel(int layout) { if (layout == 1) { return hc_affine_injection_norm_f32; } @@ -2241,7 +2340,7 @@ static auto hc_affine_injection_norm_kernel(int layout) { return hc_affine_injection_norm_f32; } -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul) { +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit) { const ggml_tensor * x = post->src[0]; const ggml_tensor * residual = post->src[1]; const ggml_tensor * weights = post->src[2]; @@ -2285,6 +2384,24 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; g.op = ggml_get_unary_op(unary); g.values = ggml_nelements(post); + if (emit) { + GGML_ASSERT(norm); + const uint3 cols = init_fastdiv_values(emit->cols); + if (emit->mmq) { + const auto launch_mmq = [&](auto write) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); }; + switch (emit->layout) { + case MMQ_Q8_1_DS_LAYOUT_D4: launch_mmq(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}); break; + case MMQ_Q8_1_DS_LAYOUT_DS4: launch_mmq(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}); break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: launch_mmq(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}); break; + default: GGML_ABORT("unsupported affine HC image layout"); + } + } else if (emit->q8) { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_norm_emit_q8_store{(block_q8_1 *) emit->q8, cols, emit->padded, (half *) emit->f16, (nv_bfloat16 *) emit->bf16}); + } else { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_norm_emit_store{(half *) emit->f16, (nv_bfloat16 *) emit->bf16}); + } + return; + } if (!norm) { const size_t shared = ((256 + a.n_embd - 2)/a.n_embd + 1)*sizeof(float); const ggml_cuda_kernel_launch_params launch(dim3((g.values + 255)/256), dim3(256), shared, ctx.stream()); diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 4a06fdc8287..35ee965aa93 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -48,4 +48,15 @@ void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_te void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); +struct ggml_cuda_hc_affine_emit_data { + void * f16 = nullptr; + void * bf16 = nullptr; + void * q8 = nullptr; + void * mmq = nullptr; + int64_t cols = 1; + int64_t padded = 0; + int64_t rows = 0; + int layout = 0; +}; + +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit = nullptr); From 78b49598facd9ec7e3a9bbf2a7e68f2ae9023b3c Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 04:31:06 -0400 Subject: [PATCH 06/27] cuda : fold affine HC gates through RMS SCALE terminals Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 20 ++++++------- ggml/src/ggml-cuda/norm.cu | 53 ++++++++++++++++++++++++++++----- ggml/src/ggml-cuda/norm.cuh | 2 +- 3 files changed, 57 insertions(+), 18 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e0009d3605c..95f8bdfdab7 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5000,7 +5000,7 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g size_t bytes; if (!ggml_cuda_hc_injection_bytes(post, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } auto hc = ggml_cuda_match_hc_post_norm(graph, post_i); - if (hc.count && graph->nodes[post_i + hc.count - 1] != (hc.mul ? hc.mul : hc.norm)) { return {}; } + if (hc.count && graph->nodes[post_i + hc.count - 1] != (hc.scale ? hc.scale : hc.mul ? hc.mul : hc.norm)) { return {}; } if (hc.count && hc.norm->ne[0] != x->ne[0]) { if (post_comb && bytes/sizeof(float) <= 256) { return {}; } if (hc.norm->ne[0] < x->ne[0] && hc.norm->ne[0] > WARP_SIZE/2 && hc.norm->ne[0] < WARP_SIZE) { return {}; } @@ -5016,7 +5016,7 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g std::vector outputs; for (int j = i; j < end; ++j) { ops.push_back(graph->nodes[j]->op); - if (!hc.mul || graph->nodes[j] != hc.norm) { outputs.push_back(j); } + if ((!hc.mul && !hc.scale) || graph->nodes[j] != hc.norm) { outputs.push_back(j); } } if (!ggml_can_fuse_subgraph(graph, i, end - i, ops.data(), outputs.data(), outputs.size())) { return {}; } if (allocated) { @@ -5031,7 +5031,7 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g std::vector writes = {affine.mul, affine.add, affine.unary}; if (affine.post) { writes.push_back(affine.post); } writes.push_back(post); - if (hc.count) { writes.push_back(hc.mul ? hc.mul : hc.norm); } + if (hc.count) { writes.push_back(hc.scale ? hc.scale : hc.mul ? hc.mul : hc.norm); } if (comb) { writes.push_back(comb); } std::vector rb(reads.size()), re(reads.size()), wb(writes.size()), we(writes.size()); for (size_t j = 0; j < reads.size(); ++j) { @@ -5867,7 +5867,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph GGML_ASSERT(ok); } const auto & affine = injection.affine; - ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, injection.post, injection.hc.norm, injection.hc.mul); + ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, injection.post, injection.hc.norm, injection.hc.mul, nullptr, injection.hc.scale); return injection.count - 1; } } @@ -6181,8 +6181,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (affine_match && affine_injection.hc.norm && prepared_hc_post >= 0) { const auto fits = [&](const auto & emit) { return !emit.norm || (emit.post == affine_injection.post && emit.norm == affine_injection.hc.norm && - emit.mul == affine_injection.hc.mul && !emit.add && !emit.scale && - emit.dst == (emit.mul ? emit.mul : emit.norm) && emit.last == i + affine_injection.count - 1); + emit.mul == affine_injection.hc.mul && !emit.add && emit.scale == affine_injection.hc.scale && + emit.dst == (emit.scale ? emit.scale : emit.mul ? emit.mul : emit.norm) && emit.last == i + affine_injection.count - 1); }; if ((norm_emits.empty() || fits(norm_emits[prepared_hc_post])) && (q8_emits.empty() || fits(q8_emits[prepared_hc_post])) && @@ -6203,7 +6203,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } const auto & affine = affine_injection.affine; ggml_cuda_op_hc_affine_injection(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post ? affine.post : affine.unary, - affine_injection.post, affine_injection.hc.norm, affine_injection.hc.mul, &data); + affine_injection.post, affine_injection.hc.norm, affine_injection.hc.mul, &data, affine_injection.hc.scale); }; if (!mmq_emits.empty() && mmq_emits[emit_i].post && (q8_emits.empty() || !q8_emits[emit_i].norm)) { const auto & emit = mmq_emits[emit_i]; @@ -6251,7 +6251,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud bf16 = image(typed.bf16); } if (emit.post) { - if (emit.scale) { + if (emit.scale && !affine_emit) { ggml_cuda_op_dsv4_hc_post(*cuda_ctx, emit.post); ggml_cuda_op_rms_norm_emit_q8(*cuda_ctx, emit.norm, nullptr, nullptr, emit.scale, f16, bf16, q8_inputs[g].get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING)); @@ -6279,7 +6279,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud return shared_inputs[g].get(); }; if (emit.post) { - if (emit.scale) { + if (emit.scale && !affine_emit) { ggml_cuda_op_hc_post_norm_scale(*cuda_ctx, emit.post, emit.norm, emit.scale, image(emit.f16), image(emit.bf16)); } else { if (affine_emit) { @@ -6556,7 +6556,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph const auto affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i, false); if (affine_injection.count) { ggml_tensor * terminal = cgraph->nodes[i + affine_injection.count - 1]; - ggml_tensor * unused_norm = affine_injection.hc.mul ? affine_injection.hc.norm : nullptr; + ggml_tensor * unused_norm = affine_injection.hc.mul || affine_injection.hc.scale ? affine_injection.hc.norm : nullptr; for (int j = i; j < i + affine_injection.count - 1; ++j) { if (cgraph->nodes[j] != unused_norm) { params->add_alloc_dep(params->user_data, cgraph->nodes[j], terminal); } for (int k = 0; k < GGML_MAX_SRC; ++k) { diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 985d34ba4ff..c5f00f40ebd 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -2307,6 +2307,16 @@ static __device__ __forceinline__ void hc_injection_norm_emit_impl(hc_post_norm_ } +struct ggml_cuda_norm_emit_q8_scale_store { + static constexpr int width = 1; + ggml_cuda_norm_emit_q8_store write; + float scale; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + write(dst, base, col, scale * value); + } +}; + template static __global__ __launch_bounds__(Block) void hc_affine_injection_emit_f32(hc_post_norm_data a, hc_affine_injection_data g, Write write) { hc_injection_norm_emit_impl(a, g, write); @@ -2319,7 +2329,7 @@ static auto hc_affine_injection_emit_kernel(int layout) { return hc_affine_injection_emit_f32; } -template +template static void ggml_cuda_hc_affine_emit_launch(ggml_backend_cuda_context & ctx, hc_post_norm_data a, hc_affine_injection_data g, ggml_tensor * norm, bool comb, bool multiply, const Write write) { const int block_size = a.ncols < 1024 ? 256 : 1024; @@ -2327,10 +2337,17 @@ static void ggml_cuda_hc_affine_emit_launch(ggml_backend_cuda_context & ctx, hc_ const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), shared, ctx.stream()); const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; - auto kernel = block_size == 256 - ? (comb ? (multiply ? hc_affine_injection_emit_kernel<256, true, true, Write>(layout) : hc_affine_injection_emit_kernel<256, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<256, false, true, Write>(layout) : hc_affine_injection_emit_kernel<256, false, false, Write>(layout))) - : (comb ? (multiply ? hc_affine_injection_emit_kernel<1024, true, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<1024, false, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, false, false, Write>(layout))); - ggml_cuda_kernel_launch(kernel, launch, a, g, write); + if constexpr (has_multiply) { + auto kernel = block_size == 256 + ? (comb ? (multiply ? hc_affine_injection_emit_kernel<256, true, true, Write>(layout) : hc_affine_injection_emit_kernel<256, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<256, false, true, Write>(layout) : hc_affine_injection_emit_kernel<256, false, false, Write>(layout))) + : (comb ? (multiply ? hc_affine_injection_emit_kernel<1024, true, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, true, false, Write>(layout)) : (multiply ? hc_affine_injection_emit_kernel<1024, false, true, Write>(layout) : hc_affine_injection_emit_kernel<1024, false, false, Write>(layout))); + ggml_cuda_kernel_launch(kernel, launch, a, g, write); + } else { + auto kernel = block_size == 256 + ? (comb ? hc_affine_injection_emit_kernel<256, true, false, Write>(layout) : hc_affine_injection_emit_kernel<256, false, false, Write>(layout)) + : (comb ? hc_affine_injection_emit_kernel<1024, true, false, Write>(layout) : hc_affine_injection_emit_kernel<1024, false, false, Write>(layout)); + ggml_cuda_kernel_launch(kernel, launch, a, g, write); + } } template @@ -2340,7 +2357,7 @@ static auto hc_affine_injection_norm_kernel(int layout) { return hc_affine_injection_norm_f32; } -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit) { +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit, ggml_tensor * scale) { const ggml_tensor * x = post->src[0]; const ggml_tensor * residual = post->src[1]; const ggml_tensor * weights = post->src[2]; @@ -2353,7 +2370,7 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens a.comb = comb ? (const float *) comb->data : nullptr; a.mul = gamma ? (const float *) gamma->data : nullptr; a.post_dst = (float *) post->data; - a.dst = norm ? (float *) (mul ? mul : norm)->data : nullptr; + a.dst = norm ? (float *) (scale ? scale : mul ? mul : norm)->data : nullptr; a.n_embd = x->ne[0]; a.hc = residual->ne[1]; a.ncols = norm ? norm->ne[0] : 256; @@ -2384,6 +2401,28 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens g.bias1 = g.last ? ggml_get_op_params_f32(last, 1) : 0.0f; g.op = ggml_get_unary_op(unary); g.values = ggml_nelements(post); + if (scale) { + GGML_ASSERT(norm && !mul); + const float factor = ggml_get_op_params_f32(scale, 0); + void * f16 = emit ? emit->f16 : nullptr; + void * bf16 = emit ? emit->bf16 : nullptr; + if (emit && emit->mmq) { + const uint3 cols = init_fastdiv_values(emit->cols); + const auto launch = [&](auto write) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); }; + switch (emit->layout) { + case MMQ_Q8_1_DS_LAYOUT_D4: launch(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}); break; + case MMQ_Q8_1_DS_LAYOUT_DS4: launch(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}); break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: launch(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}); break; + default: GGML_ABORT("unsupported affine HC image layout"); + } + } else if (emit && emit->q8) { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, ggml_cuda_norm_emit_q8_scale_store{ + {(block_q8_1 *) emit->q8, init_fastdiv_values(emit->cols), emit->padded, (half *) f16, (nv_bfloat16 *) bf16}, factor}); + } else { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, ggml_cuda_norm_scale_store{factor, (half *) f16, (nv_bfloat16 *) bf16}); + } + return; + } if (emit) { GGML_ASSERT(norm); const uint3 cols = init_fastdiv_values(emit->cols); diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 35ee965aa93..2c80cc8d429 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -59,4 +59,4 @@ struct ggml_cuda_hc_affine_emit_data { int layout = 0; }; -void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit = nullptr); +void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit = nullptr, ggml_tensor * scale = nullptr); From d355462afd56638506f9c405a5c9ee9d23830474 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 08:09:17 -0400 Subject: [PATCH 07/27] cuda : fuse cuBLAS output conversion with gated HC reduction Assisted-by: Codex --- ggml/src/ggml-cuda/dsv4-hc.cu | 88 +++++++++++++++++++++++++++++++-- ggml/src/ggml-cuda/dsv4-hc.cuh | 2 + ggml/src/ggml-cuda/ggml-cuda.cu | 75 ++++++++++++++++++---------- tests/test-backend-ops.cpp | 51 +++++++++++++++++-- 4 files changed, 180 insertions(+), 36 deletions(-) diff --git a/ggml/src/ggml-cuda/dsv4-hc.cu b/ggml/src/ggml-cuda/dsv4-hc.cu index ca1d2dc8a48..c55c1cebe91 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cu +++ b/ggml/src/ggml-cuda/dsv4-hc.cu @@ -1,5 +1,6 @@ #include "common.cuh" #include "dsv4-hc.cuh" +#include "convert.cuh" static constexpr int DSV4_HC = 4; @@ -100,10 +101,26 @@ static __global__ void dsv4_hc_comb_f32( } } -template -static __global__ void dsv4_hc_pre_f32( +struct ggml_cuda_hc_gate_f32 { + const float * data; + __device__ __forceinline__ float operator()(int64_t i) const { return data[i]; } +}; + +template +struct ggml_cuda_hc_gate_convert { + const T * data; + float * raw; + __device__ __forceinline__ float operator()(int64_t i) const { + const float value = ggml_cuda_cast(data[i]); + raw[i] = value; + return value; + } +}; + +template +static __device__ __forceinline__ void dsv4_hc_pre_impl( const float * x, - const float * weights, + const Gate weights, float * dst, int64_t n_embd, int64_t hc, @@ -135,9 +152,9 @@ static __global__ void dsv4_hc_pre_f32( const float xv = x[i0*sx0 + ih*sx1 + it*sx2]; float wv; if constexpr (gated) { - wv = 1.0f / (1.0f + expf(-weights[i0*sw0 + ih*sw1 + it*sw2])); + wv = 1.0f / (1.0f + expf(-weights(i0*sw0 + ih*sw1 + it*sw2))); } else { - wv = weights[ih*sw0 + it*sw1]; + wv = weights(ih*sw0 + it*sw1); } sum += xv * wv; } @@ -145,6 +162,40 @@ static __global__ void dsv4_hc_pre_f32( dst[i0*sd0 + it*sd1] = scale * sum; } +template +static __global__ void dsv4_hc_pre_f32( + const float * x, + const float * weights, + float * dst, + int64_t n_embd, + int64_t hc, + int64_t n_tokens, + int64_t sx0, + int64_t sx1, + int64_t sx2, + int64_t sw0, + int64_t sw1, + int64_t sw2, + int64_t sd0, + int64_t sd1, + float scale) { + dsv4_hc_pre_impl( + x, ggml_cuda_hc_gate_f32{weights}, dst, n_embd, hc, n_tokens, + sx0, sx1, sx2, sw0, sw1, sw2, sd0, sd1, scale); +} + +template +static __global__ void dsv4_hc_pre_convert( + const float * x, const void * weights, float * raw, float * dst, + int64_t n_embd, int64_t hc, int64_t n_tokens, + int64_t sx0, int64_t sx1, int64_t sx2, + int64_t sw0, int64_t sw1, int64_t sw2, + int64_t sd0, int64_t sd1, float scale) { + dsv4_hc_pre_impl( + x, ggml_cuda_hc_gate_convert{(const T *) weights, raw}, dst, n_embd, hc, n_tokens, + sx0, sx1, sx2, sw0, sw1, sw2, sd0, sd1, scale); +} + template static __global__ void dsv4_hc_post_f32( const float * x, @@ -272,6 +323,33 @@ void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst scale); } +void ggml_cuda_op_dsv4_hc_pre_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * weights, ggml_tensor * mm, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * view = dst->src[1]; + GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_BF16); + GGML_ASSERT(x->type == GGML_TYPE_F32 && mm->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_get_op_params_i32(dst, 1) && ggml_is_contiguous(mm) && ggml_is_contiguous(view)); + + const int64_t nr = x->ne[0]*x->ne[2]; + const dim3 block_dims(256, 1, 1); + const dim3 grid_dims((nr + block_dims.x - 1)/block_dims.x, 1, 1); + const ggml_cuda_kernel_launch_params params(grid_dims, block_dims, 0, ctx.stream()); + const auto launch = [&](auto tag) { + using T = decltype(tag); + ggml_cuda_kernel_launch(dsv4_hc_pre_convert, params, + (const float *) x->data, weights, (float *) mm->data, (float *) dst->data, + x->ne[0], x->ne[1], x->ne[2], + x->nb[0]/sizeof(float), x->nb[1]/sizeof(float), x->nb[2]/sizeof(float), + view->nb[0]/sizeof(float), view->nb[1]/sizeof(float), view->nb[2]/sizeof(float), + dst->nb[0]/sizeof(float), dst->nb[1]/sizeof(float), ggml_get_op_params_f32(dst, 0)); + }; + if (type == GGML_TYPE_F16) { + launch(half{}); + } else { + launch(nv_bfloat16{}); + } +} + void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * x = dst->src[0]; const ggml_tensor * residual = dst->src[1]; diff --git a/ggml/src/ggml-cuda/dsv4-hc.cuh b/ggml/src/ggml-cuda/dsv4-hc.cuh index 2379aaefb41..aa3cbae43eb 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cuh +++ b/ggml/src/ggml-cuda/dsv4-hc.cuh @@ -4,3 +4,5 @@ void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_dsv4_hc_pre_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * weights, ggml_tensor * mm, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 95f8bdfdab7..10d10389224 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1401,8 +1401,18 @@ struct batched_mul_mat_traits { static inline auto convert_nc(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } }; +static bool ggml_cuda_cublas_prefer_f32_output(ggml_type type, int cc) { + if (type == GGML_TYPE_F16) { + return cc == GGML_CUDA_CC_VOLTA || GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_CDNA(cc); + } + if (type == GGML_TYPE_BF16) { + return !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc); + } + return false; +} + template -static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr) { +static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr) { using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; @@ -1512,12 +1522,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const const void * beta = traits::get_beta(); const int cc = ggml_cuda_info().devices[ctx.device].cc; - bool prefer_f32_output = false; - if (compute_type == GGML_TYPE_F16) { - prefer_f32_output = cc == GGML_CUDA_CC_VOLTA || GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_CDNA(cc); - } else if (compute_type == GGML_TYPE_BF16) { - prefer_f32_output = !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc); - } + const bool prefer_f32_output = ggml_cuda_cublas_prefer_f32_output(compute_type, cc); if (prefer_f32_output) { dst_ptr = (char *) dst_ddf; @@ -1619,8 +1624,12 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const // Convert output back to F32 if needed if (cu_data_type != CUDA_R_32F) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); - to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); + if (hc_dst) { + ggml_cuda_op_dsv4_hc_pre_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, hc_dst); + } else { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); + to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); + } } } @@ -1666,17 +1675,17 @@ static ggml_type ggml_cuda_mul_mat_cublas_compute_type(int cc, const ggml_tensor return compute_type; } -static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr) { +static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr) { const ggml_type compute_type = ggml_cuda_mul_mat_cublas_compute_type(ggml_cuda_info().devices[ctx.device].cc, src0, src1, dst, true); switch (compute_type) { case GGML_TYPE_F32: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); break; case GGML_TYPE_BF16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); break; case GGML_TYPE_F16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); break; default: GGML_ABORT("fatal error"); @@ -1847,7 +1856,7 @@ static ggml_cuda_mm_kernel ggml_cuda_mul_mat_kernel(const ggml_tensor * src0, co // 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. // Therefore, in such cases use cuBLAS. - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE + const bool bad_padding_clear = src0->buffer && 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) { return GGML_CUDA_MM_CUBLAS; @@ -4679,12 +4688,13 @@ struct ggml_cuda_hc_up_match { ggml_tensor * mm = nullptr; ggml_tensor * dst = nullptr; int count = 0; + bool convert = false; }; -static ggml_cuda_hc_up_match ggml_cuda_match_hc_up_shape(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i) { +static ggml_cuda_hc_up_match ggml_cuda_match_hc_up_shape(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool convert = false) { ggml_tensor * mm = graph->nodes[i]; if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || - mm->src[1]->type != GGML_TYPE_F32 || mm->ne[2] != 1 || mm->ne[3] != 1 || mm->ne[1] > MMVF_MAX_BATCH_SIZE || + mm->src[1]->type != GGML_TYPE_F32 || mm->ne[2] != 1 || mm->ne[3] != 1 || (!convert && mm->ne[1] > MMVF_MAX_BATCH_SIZE) || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } @@ -4701,20 +4711,26 @@ static ggml_cuda_hc_up_match ggml_cuda_match_hc_up_shape(ggml_backend_cuda_conte !ggml_are_same_shape(view, norm) || norm->ne[3] != 1 || mm->ne[0]/norm->ne[0] != norm->ne[1] || mm->ne[0] % norm->ne[0] || mm->ne[1] != norm->ne[2] || dst->ne[0] != norm->ne[0] || dst->ne[1] != norm->ne[2] || dst->ne[2] != 1 || dst->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || input->ne[2] != 1 || input->ne[3] != 1 || - (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16 && weight->type != GGML_TYPE_BF16)) { + (!convert && weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16 && weight->type != GGML_TYPE_BF16)) { return {}; } - if (!ggml_cuda_should_use_mmvf(weight->type, ggml_cuda_info().devices[ctx.device].cc, weight->ne, weight->nb, input->ne[1]) || + if (convert) { + const int cc = ggml_cuda_info().devices[ctx.device].cc; + if (ggml_cuda_mul_mat_kernel(weight, input, mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return {}; } + const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(cc, weight, input, mm, false); + if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || ggml_cuda_cublas_prefer_f32_output(type, cc)) { return {}; } + } else if (!ggml_cuda_should_use_mmvf(weight->type, ggml_cuda_info().devices[ctx.device].cc, weight->ne, weight->nb, input->ne[1]) || !ggml_cuda_should_fuse_hc_up(weight, norm, ctx.device)) { return {}; } const int count = last - i + 1; ggml_op ops[4]; int outputs[4]; for (int j = 0; j < count; ++j) { ops[j] = graph->nodes[i + j]->op; outputs[j] = i + j; } if (!ggml_can_fuse_subgraph(graph, i, count, ops, outputs, count)) { return {}; } - return {mm, dst, count}; + return {mm, dst, count, convert}; } static ggml_cuda_hc_up_match ggml_cuda_match_hc_up(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i) { - const auto match = ggml_cuda_match_hc_up_shape(ctx, graph, i); + auto match = ggml_cuda_match_hc_up_shape(ctx, graph, i); + if (!match.count) { match = ggml_cuda_match_hc_up_shape(ctx, graph, i, true); } if (!match.count) { return {}; } ggml_tensor * mm = match.mm; ggml_tensor * dst = match.dst; @@ -4726,10 +4742,10 @@ static ggml_cuda_hc_up_match ggml_cuda_match_hc_up(ggml_backend_cuda_context & c for (const ggml_tensor * tensor : tensors) { uintptr_t begin, end; if (!ggml_cuda_hc_up_range(tensor, ctx.device, begin, end) || tensor->nb[0] != ggml_type_size(tensor->type) || - (end - begin)/ggml_type_size(tensor->type) > INT_MAX || begin % (2*ggml_type_size(tensor->type))) { return {}; } + (end - begin)/ggml_type_size(tensor->type) > INT_MAX || begin % (match.convert ? sizeof(float) : 2*ggml_type_size(tensor->type))) { return {}; } for (int d = 1; d < GGML_MAX_DIMS; ++d) { if (tensor->ne[d] > INT_MAX || tensor->nb[d]/ggml_type_size(tensor->type) > INT_MAX || - (tensor != mm && tensor->nb[d] % (2*ggml_type_size(tensor->type)))) { return {}; } + (tensor != mm && tensor->nb[d] % ((match.convert ? 1 : 2)*ggml_type_size(tensor->type)))) { return {}; } } } if (!(mm->flags & GGML_TENSOR_FLAG_COMPUTE) || !(dst->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_is_contiguous(mm) || @@ -5059,7 +5075,7 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, - const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr) { + const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr, const void * prepared_src1 = nullptr) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (disable_fusion) { @@ -5850,7 +5866,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto match = ggml_cuda_match_hc_up(*cuda_ctx, cgraph, i); if (match.count) { - ggml_cuda_mul_mat_vec_f_hc_pre(*cuda_ctx, match.mm, match.dst); + if (match.convert) { + ggml_cuda_mul_mat_cublas(*cuda_ctx, match.mm->src[0], match.mm->src[1], match.mm, prepared_src1, match.dst); + } else { + ggml_cuda_mul_mat_vec_f_hc_pre(*cuda_ctx, match.mm, match.dst); + } return match.count - 1; } } @@ -6187,7 +6207,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if ((norm_emits.empty() || fits(norm_emits[prepared_hc_post])) && (q8_emits.empty() || fits(q8_emits[prepared_hc_post])) && (mmq_emits.empty() || fits(mmq_emits[prepared_hc_post]))) { - const int skipped = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match, &affine_emit); + const int skipped = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match, &affine_emit, prepared_src1); if (skipped) { i += skipped; continue; @@ -6337,7 +6357,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } } - int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match); + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match, nullptr, prepared_src1); if (nodes_to_skip != 0) { #ifdef GGML_CUDA_DEBUG @@ -6547,7 +6567,8 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { - const auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); + auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); + if (!hc_up.count) { hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i, true); } if (hc_up.count) { params->add_alloc_dep(params->user_data, hc_up.mm->src[0], hc_up.dst); params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 081fb23ff12..675cfe82a5e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4319,6 +4319,11 @@ struct test_dsv4_hc_pre : public test_dsv4_hc { const int64_t n_hc; const int64_t n_tokens; const bool gated; + const ggml_type projection_type; + const bool computed; + ggml_tensor * projection = nullptr; + ggml_tensor * projection_weights = nullptr; + ggml_tensor * projection_input = nullptr; std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); @@ -4326,18 +4331,53 @@ struct test_dsv4_hc_pre : public test_dsv4_hc { } std::string vars() override { - return VARS_TO_STR4(n_embd, n_hc, n_tokens, gated); + return VARS_TO_STR4(n_embd, n_hc, n_tokens, gated) + (projection_type == GGML_TYPE_COUNT ? "" : "," + VARS_TO_STR2(projection_type, computed)); } - test_dsv4_hc_pre(int64_t n_embd = 31, int64_t n_hc = 4, int64_t n_tokens = 17, bool gated = false) - : n_embd(n_embd), n_hc(n_hc), n_tokens(n_tokens), gated(gated) {} + bool run_whole_graph() override { return projection_type != GGML_TYPE_COUNT; } + std::vector fusion_test_nodes() override { return projection ? std::vector{projection, out} : std::vector{}; } + + test_dsv4_hc_pre(int64_t n_embd = 31, int64_t n_hc = 4, int64_t n_tokens = 17, bool gated = false, ggml_type projection_type = GGML_TYPE_COUNT, bool computed = false) + : n_embd(n_embd), n_hc(n_hc), n_tokens(n_tokens), gated(gated), projection_type(projection_type), computed(computed) {} + + void initialize_tensors(ggml_context * ctx) override { + test_dsv4_hc::initialize_tensors(ctx); + if (!projection) { return; } + // signed powers of two keep the projection sums exact in F16 + for (ggml_tensor * t : {projection_weights, projection_input}) { + std::vector data(ggml_nelements(t)); + for (int64_t row = 0; row < t->ne[1]; ++row) { + for (int64_t col = 0; col < t->ne[0]; ++col) { + data[row*t->ne[0] + col] = ((row*13 + col*7) % 11 < 5 ? -1.0f : 1.0f)/32.0f; + } + } + std::vector bytes(ggml_nbytes(t)); + ggml_quantize_chunk(t->type, data.data(), bytes.data(), 0, t->ne[1], t->ne[0], nullptr); + ggml_backend_tensor_set(t, bytes.data(), 0, bytes.size()); + } + } ggml_tensor * build_graph(ggml_context * ctx) override { + projection = nullptr; + projection_weights = nullptr; + projection_input = nullptr; ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_hc, n_tokens); ggml_set_name(x, "x"); if (gated) { - ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_hc, n_tokens); + ggml_tensor * gate; + if (projection_type != GGML_TYPE_COUNT) { + ggml_tensor * weights = ggml_new_tensor_2d(ctx, projection_type, 127, n_embd*n_hc); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 127, n_tokens); + projection_weights = weights; + projection_input = input; + if (computed) { weights = ggml_cont(ctx, weights); } + projection = ggml_mul_mat(ctx, weights, input); + ggml_set_output(projection); + gate = ggml_reshape_3d(ctx, projection, n_embd, n_hc, n_tokens); + } else { + gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_hc, n_tokens); + } ggml_set_name(gate, "gate"); out = ggml_dsv4_hc_pre_gated(ctx, x, gate, 1.0f/n_hc); @@ -9253,6 +9293,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, 21)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 4, 17, true)); test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, 21, true)); + test_cases.emplace_back(new test_dsv4_hc_pre(512, 4, 32, true, GGML_TYPE_F16)); + test_cases.emplace_back(new test_dsv4_hc_pre(257, 3, 17, true, GGML_TYPE_F16, true)); + test_cases.emplace_back(new test_dsv4_hc_pre(257, 3, 17, true, GGML_TYPE_BF16)); for (int64_t n_hc : {1, 2, 3, 5, 8, 65}) { test_cases.emplace_back(new test_dsv4_hc_pre(128, n_hc, 17)); test_cases.emplace_back(new test_dsv4_hc_pre(128, n_hc, 17, true)); From 6c02602660e9d799e9e56287746c0424397f0170 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 09:35:40 -0400 Subject: [PATCH 08/27] cuda : fuse cuBLAS result conversion with scalar activation Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 66 +++++++++++++++++++++++++++++---- ggml/src/ggml-cuda/unary.cu | 66 ++++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/unary.cuh | 8 ++++ tests/test-backend-ops.cpp | 63 +++++++++++++++++++++++++++++++ 4 files changed, 190 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 10d10389224..48d51754bd1 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1412,7 +1412,7 @@ static bool ggml_cuda_cublas_prefer_f32_output(ggml_type type, int cc) { } template -static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr) { +static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr) { using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; @@ -1626,6 +1626,8 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const if (cu_data_type != CUDA_R_32F) { if (hc_dst) { ggml_cuda_op_dsv4_hc_pre_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, hc_dst); + } else if (unary) { + ggml_cuda_op_scaled_unary_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, *unary); } else { const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); @@ -1675,17 +1677,17 @@ static ggml_type ggml_cuda_mul_mat_cublas_compute_type(int cc, const ggml_tensor return compute_type; } -static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr) { +static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr) { const ggml_type compute_type = ggml_cuda_mul_mat_cublas_compute_type(ggml_cuda_info().devices[ctx.device].cc, src0, src1, dst, true); switch (compute_type) { case GGML_TYPE_F32: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); break; case GGML_TYPE_BF16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); break; case GGML_TYPE_F16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); break; default: GGML_ABORT("fatal error"); @@ -4390,7 +4392,7 @@ struct ggml_cuda_scaled_unary_match { int count = 0; }; -static ggml_cuda_scaled_unary_match ggml_cuda_match_scaled_unary(ggml_cgraph * cgraph, int i) { +static ggml_cuda_scaled_unary_match ggml_cuda_match_scaled_unary(ggml_cgraph * cgraph, int i, bool plain = false) { ggml_cuda_scaled_unary_match match; ggml_tensor * first = cgraph->nodes[i]; const bool before = first->op == GGML_OP_SCALE; @@ -4399,7 +4401,7 @@ static ggml_cuda_scaled_unary_match ggml_cuda_match_scaled_unary(ggml_cgraph * c (ggml_get_unary_op(cgraph->nodes[u]) != GGML_UNARY_OP_SILU && ggml_get_unary_op(cgraph->nodes[u]) != GGML_UNARY_OP_SIGMOID)) { return match; } - if (before && (ggml_cuda_can_fuse(cgraph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || + if ((before || plain) && (ggml_cuda_can_fuse(cgraph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || ggml_cuda_can_fuse(cgraph, u, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }))) { return match; } @@ -4415,6 +4417,7 @@ static ggml_cuda_scaled_unary_match ggml_cuda_match_scaled_unary(ggml_cgraph * c } else if (!before && after && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_UNARY, GGML_OP_SCALE }, { i, i + 1 })) { count = 2; } + if (!before && !after && plain && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_UNARY }, { i })) { count = 1; } if (!count || !first->src[0] || (first->type != GGML_TYPE_F32 && first->type != GGML_TYPE_BF16)) { return match; } @@ -4771,6 +4774,41 @@ static ggml_cuda_hc_up_match ggml_cuda_match_hc_up(ggml_backend_cuda_context & c return match; } +struct ggml_cuda_cublas_unary_match { + ggml_tensor * mm = nullptr; + ggml_cuda_scaled_unary_match unary; + int count = 0; +}; + +static ggml_cuda_cublas_unary_match ggml_cuda_match_cublas_unary(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * mm = graph->nodes[i]; + if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || + mm->src[1]->type != GGML_TYPE_F32 || !ggml_is_contiguous(mm) || i + 1 >= graph->n_nodes || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } + const auto unary = ggml_cuda_match_scaled_unary(graph, i + 1, true); + if (!unary.count || unary.first->src[0] != mm || ggml_cuda_match_hc_injection(ctx, graph, i + 1, false).count) { return {}; } + const int cc = ggml_cuda_info().devices[ctx.device].cc; + if (ggml_cuda_mul_mat_kernel(mm->src[0], mm->src[1], mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return {}; } + const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(cc, mm->src[0], mm->src[1], mm, false); + if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || ggml_cuda_cublas_prefer_f32_output(type, cc)) { return {}; } + const int count = unary.count + 1; + ggml_op ops[4]; int outputs[4]; + for (int j = 0; j < count; ++j) { ops[j] = graph->nodes[i + j]->op; outputs[j] = i + j; } + if (!ggml_can_fuse_subgraph(graph, i, count, ops, outputs, count)) { return {}; } + if (allocated) { + if (!ggml_cuda_scaled_unary_memory_ok(unary, ctx.device)) { return {}; } + uintptr_t read_begin[2], read_end[2]; + for (int j = 0; j < 2; ++j) { + if (!ggml_cuda_hc_up_range(mm->src[j], ctx.device, read_begin[j], read_end[j])) { return {}; } + } + for (const ggml_tensor * tensor : {mm, unary.first, unary.unary, unary.last}) { + uintptr_t begin, end; + if (!(tensor->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_cuda_hc_up_range(tensor, ctx.device, begin, end)) { return {}; } + for (int j = 0; j < 2; ++j) { if (begin < read_end[j] && read_begin[j] < end) { return {}; } } + } + } + return {mm, unary, count}; +} + static bool ggml_cuda_affine_unary_bytes(const ggml_tensor * tensor, size_t & bytes) { for (int d = 0; d < GGML_MAX_DIMS; ++d) { if (tensor->ne[d] <= 0) { return false; } @@ -5863,6 +5901,15 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_cublas_unary(*cuda_ctx, cgraph, i); + if (match.count) { + const ggml_cuda_scaled_unary_args args = {match.unary.first, match.unary.unary, match.unary.last}; + ggml_cuda_mul_mat_cublas(*cuda_ctx, match.mm->src[0], match.mm->src[1], match.mm, prepared_src1, nullptr, &args); + return match.count - 1; + } + } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto match = ggml_cuda_match_hc_up(*cuda_ctx, cgraph, i); if (match.count) { @@ -6567,6 +6614,11 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { + const auto unary = ggml_cuda_match_cublas_unary(*cuda_ctx, cgraph, i, false); + if (unary.count) { + params->add_alloc_dep(params->user_data, unary.mm->src[0], unary.unary.last); + params->add_alloc_dep(params->user_data, unary.mm->src[1], unary.unary.last); + } auto hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i); if (!hc_up.count) { hc_up = ggml_cuda_match_hc_up_shape(*cuda_ctx, cgraph, i, true); } if (hc_up.count) { diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 28954cb6970..56affe6e614 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -838,16 +838,33 @@ void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_n /* fused scale and activation */ -template -static __global__ void scaled_unary_kernel(const T * x, T * pre_dst, T * unary_dst, T * dst, int k, float scale0, float bias0, float scale1, float bias1) { +template +struct ggml_cuda_scaled_unary_load { + const T * data; + __device__ __forceinline__ T operator()(int i) const { return data[i]; } +}; + +template +struct ggml_cuda_scaled_unary_convert_load { + const T * data; + float * raw; + __device__ __forceinline__ float operator()(int i) const { + const float value = ggml_cuda_cast(data[i]); + raw[i] = value; + return value; + } +}; + +template +static __device__ __forceinline__ void scaled_unary_impl(Load x, T * pre_dst, T * unary_dst, T * dst, int k, float scale0, float bias0, float scale1, float bias1, bool before, bool after) { ggml_cuda_pdl_lc(); const int i = blockDim.x*blockIdx.x + threadIdx.x; if (i >= k) { return; } ggml_cuda_pdl_sync(); - float value = ggml_cuda_cast(x[i]); - if constexpr (before) { + float value = ggml_cuda_cast(x(i)); + if (before) { const T pre = ggml_cuda_cast(scale0 * value + bias0); pre_dst[i] = pre; value = ggml_cuda_cast(pre); @@ -855,14 +872,24 @@ static __global__ void scaled_unary_kernel(const T * x, T * pre_dst, T * unary_d const T unary = ggml_cuda_cast(op(value)); unary_dst[i] = unary; value = ggml_cuda_cast(unary); - if constexpr (after) { + if (after) { value = scale1 * value + bias1; } - if constexpr (after) { + if (after) { dst[i] = ggml_cuda_cast(value); } } +template +static __global__ void scaled_unary_kernel(const T * x, T * pre_dst, T * unary_dst, T * dst, int k, float scale0, float bias0, float scale1, float bias1) { + scaled_unary_impl(ggml_cuda_scaled_unary_load{x}, pre_dst, unary_dst, dst, k, scale0, bias0, scale1, bias1, before, after); +} + +template +static __global__ void scaled_unary_convert_kernel(const void * x, float * raw, float * pre_dst, float * unary_dst, float * dst, int k, float scale0, float bias0, float scale1, float bias1, bool before, bool after) { + scaled_unary_impl(ggml_cuda_scaled_unary_convert_load{(const T *) x, raw}, pre_dst, unary_dst, dst, k, scale0, bias0, scale1, bias1, before, after); +} + template static void scaled_unary_cuda(const ggml_tensor * first, const ggml_tensor * unary, const ggml_tensor * last, bool before, bool after, cudaStream_t stream) { const int k = ggml_nelements(last); @@ -907,3 +934,30 @@ void ggml_cuda_op_scaled_unary(ggml_backend_cuda_context & ctx, ggml_tensor * fi GGML_ABORT("Unsupported scaled unary op"); } } + +template +static void scaled_unary_convert_cuda(ggml_backend_cuda_context & ctx, const void * src, ggml_tensor * mm, const ggml_cuda_scaled_unary_args & args) { + const bool before = args.first->op == GGML_OP_SCALE; + const bool after = args.last->op == GGML_OP_SCALE; + const int k = ggml_nelements(mm); + const ggml_cuda_kernel_launch_params launch((k + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(scaled_unary_convert_kernel, launch, src, (float *) mm->data, + before ? (float *) args.first->data : nullptr, (float *) args.unary->data, (float *) args.last->data, k, + before ? ggml_get_op_params_f32(args.first, 0) : 1.0f, before ? ggml_get_op_params_f32(args.first, 1) : 0.0f, + after ? ggml_get_op_params_f32(args.last, 0) : 1.0f, after ? ggml_get_op_params_f32(args.last, 1) : 0.0f, before, after); +} + +void ggml_cuda_op_scaled_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_scaled_unary_args & args) { + GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_BF16); + GGML_ASSERT(mm->type == GGML_TYPE_F32 && args.first->type == GGML_TYPE_F32 && args.unary->type == GGML_TYPE_F32 && args.last->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(mm) && ggml_are_same_shape(mm, args.last)); + const auto launch = [&](auto tag) { + using T = decltype(tag); + switch (ggml_get_unary_op(args.unary)) { + case GGML_UNARY_OP_SILU: scaled_unary_convert_cuda(ctx, src, mm, args); break; + case GGML_UNARY_OP_SIGMOID: scaled_unary_convert_cuda(ctx, src, mm, args); break; + default: GGML_ABORT("Unsupported scaled unary op"); + } + }; + if (type == GGML_TYPE_F16) { launch(half{}); } else { launch(nv_bfloat16{}); } +} diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index dafb4c46f98..fc0f124515b 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -125,3 +125,11 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl } void ggml_cuda_op_scaled_unary(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last); + +struct ggml_cuda_scaled_unary_args { + const ggml_tensor * first; + const ggml_tensor * unary; + const ggml_tensor * last; +}; + +void ggml_cuda_op_scaled_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_scaled_unary_args & args); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 675cfe82a5e..05f7c6c2441 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4148,6 +4148,61 @@ struct test_unary_mul : public test_case { } }; +struct test_mul_mat_unary : public test_case { + const ggml_type weight_type; + const ggml_unary_op op; + const bool before; + const bool after; + const bool batched; + std::vector stages; + ggml_tensor * weights = nullptr; + ggml_tensor * input = nullptr; + + test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched) + : weight_type(weight_type), op(op), before(before), after(after), batched(batched) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MUL_MAT_UNARY"; + } + + std::string vars() override { return VARS_TO_STR5(weight_type, op, before, after, batched); } + bool run_whole_graph() override { return true; } + std::vector fusion_test_nodes() override { return stages; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + weights = ggml_new_tensor_4d(ctx, weight_type, 127, 129, batched ? 2 : 1, 1); + input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 127, 17, batched ? 2 : 1, 1); + ggml_tensor * out = ggml_mul_mat(ctx, weights, input); + stages = {out}; + if (before) { out = ggml_scale_bias(ctx, out, 0.75f, 0.125f); stages.push_back(out); } + out = ggml_unary(ctx, out, op); + stages.push_back(out); + if (after) { out = ggml_scale_bias(ctx, out, 1.5f, -0.025f); stages.push_back(out); } + for (ggml_tensor * stage : stages) { ggml_set_output(stage); } + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { init_tensor_uniform(t); } + // Exact projection sums keep the original backend tolerance valid. + for (ggml_tensor * t : {weights, input}) { + std::vector values(ggml_nelements(t)); + for (size_t i = 0; i < values.size(); ++i) { values[i] = ((i*13 + i/127)%7 < 3 ? -1.0f : 1.0f)/32.0f; } + if (t->type == GGML_TYPE_F16) { + std::vector typed(values.size()); + ggml_fp32_to_fp16_row(values.data(), typed.data(), typed.size()); + ggml_backend_tensor_set(t, typed.data(), 0, ggml_nbytes(t)); + } else if (t->type == GGML_TYPE_BF16) { + std::vector typed(values.size()); + ggml_fp32_to_bf16_row(values.data(), typed.data(), typed.size()); + ggml_backend_tensor_set(t, typed.data(), 0, ggml_nbytes(t)); + } else { ggml_backend_tensor_set(t, values.data(), 0, ggml_nbytes(t)); } + } + } +}; + // SNAKE activation fusion: y = x + sin(a*x)^2 * inv_b // CUDA backend matches the naive 5-op chain (mul, sin, sqr, mul, add) // and dispatches a single fused kernel. @@ -9277,6 +9332,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_snake_fuse(type, { 64, 32, 2, 3})); // ne[2] > 1 and ne[3] > 1 } + for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32}) { + for (ggml_unary_op op : {GGML_UNARY_OP_SILU, GGML_UNARY_OP_SIGMOID}) { + test_cases.emplace_back(new test_mul_mat_unary(type, op, true, true, false)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true)); + } + } + test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); From 3df94da2df30ade9caedeca051b0b93592755e1e Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 11:34:20 -0400 Subject: [PATCH 09/27] cuda : fold cuBLAS conversion into tensor affine activation Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 82 ++++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/unary.cu | 47 +++++++++++++++++-- ggml/src/ggml-cuda/unary.cuh | 9 ++++ tests/test-backend-ops.cpp | 17 +++++-- 4 files changed, 143 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 48d51754bd1..4fee77cd3ce 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1412,7 +1412,7 @@ static bool ggml_cuda_cublas_prefer_f32_output(ggml_type type, int cc) { } template -static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr) { +static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr, const ggml_cuda_affine_unary_ops * affine = nullptr) { using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; @@ -1628,6 +1628,8 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_cuda_op_dsv4_hc_pre_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, hc_dst); } else if (unary) { ggml_cuda_op_scaled_unary_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, *unary); + } else if (affine) { + ggml_cuda_op_affine_unary_convert(ctx, traits::ggml_type_val, dst_temp.get(), dst, *affine); } else { const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); @@ -1677,17 +1679,17 @@ static ggml_type ggml_cuda_mul_mat_cublas_compute_type(int cc, const ggml_tensor return compute_type; } -static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr) { +static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const void * prepared_src1 = nullptr, ggml_tensor * hc_dst = nullptr, const ggml_cuda_scaled_unary_args * unary = nullptr, const ggml_cuda_affine_unary_ops * affine = nullptr) { const ggml_type compute_type = ggml_cuda_mul_mat_cublas_compute_type(ggml_cuda_info().devices[ctx.device].cc, src0, src1, dst, true); switch (compute_type) { case GGML_TYPE_F32: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary, affine); break; case GGML_TYPE_BF16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary, affine); break; case GGML_TYPE_F16: - ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary); + ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst, prepared_src1, hc_dst, unary, affine); break; default: GGML_ABORT("fatal error"); @@ -5067,12 +5069,14 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g const size_t shared = ((hc.count ? 32 : 0) + (ncols + x->ne[0] - 2)/x->ne[0] + 1)*sizeof(float); if (shared > ggml_cuda_info().devices[ctx.device].smpb) { return {}; } std::vector ops; + std::vector indices; std::vector outputs; for (int j = i; j < end; ++j) { ops.push_back(graph->nodes[j]->op); + indices.push_back(j); if ((!hc.mul && !hc.scale) || graph->nodes[j] != hc.norm) { outputs.push_back(j); } } - if (!ggml_can_fuse_subgraph(graph, i, end - i, ops.data(), outputs.data(), outputs.size())) { return {}; } + if (!ggml_can_fuse_subgraph_ext(graph, indices.data(), indices.size(), ops.data(), outputs.data(), outputs.size())) { return {}; } if (allocated) { if (hc.count && !ggml_cuda_hc_post_norm_memory_ok(graph, post_i, hc, ctx.device)) { return {}; } std::vector reads = {gate_inputs[0], gate_inputs[1], gate_inputs[2], x, residual, post_comb}; @@ -5111,6 +5115,55 @@ static ggml_cuda_hc_affine_injection_match ggml_cuda_match_hc_affine_injection(g return {affine, post, comb, hc, end - i}; } +struct ggml_cuda_cublas_affine_match { + ggml_tensor * mm = nullptr; + ggml_cuda_affine_unary_match affine; + int count = 0; +}; + +static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * mm = graph->nodes[i]; + if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || mm->src[1]->type != GGML_TYPE_F32 || + !ggml_is_contiguous(mm) || i + 1 >= graph->n_nodes || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } + int first = i + 1; + while (first < graph->n_nodes && graph->nodes[first]->op == GGML_OP_VIEW) { + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[first])) { return {}; } + uintptr_t begin, end; + if (allocated && !ggml_cuda_affine_unary_range(graph->nodes[first], ctx.device, begin, end)) { return {}; } + ++first; + } + if (first >= graph->n_nodes) { return {}; } + const auto affine = ggml_cuda_match_affine_unary(ctx, graph, first, allocated); + if (!affine.count || affine.mul->src[0] != mm || ggml_cuda_match_hc_affine_injection(ctx, graph, first, false).count) { return {}; } + if (ggml_cuda_mul_mat_kernel(mm->src[0], mm->src[1], mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return {}; } + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(cc, mm->src[0], mm->src[1], mm, false); + if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || ggml_cuda_cublas_prefer_f32_output(type, cc)) { return {}; } + const int count = first - i + affine.count; + std::vector ops; + std::vector outputs; + for (int j = i; j < i + count; ++j) { ops.push_back(graph->nodes[j]->op); outputs.push_back(j); } + if (!ggml_can_fuse_subgraph_ext(graph, outputs.data(), count, ops.data(), outputs.data(), outputs.size())) { return {}; } + if (allocated) { + uintptr_t rb[4], re[4], wb[5], we[5]; + const ggml_tensor * reads[] = {mm->src[0], mm->src[1], affine.mul->src[1], affine.add->src[1]}; + const ggml_tensor * writes[] = {mm, affine.mul, affine.add, affine.unary, affine.post}; + for (int j = 0; j < 4; ++j) { + if (j < 2 ? !ggml_cuda_hc_up_range(reads[j], ctx.device, rb[j], re[j]) : !ggml_cuda_affine_unary_range(reads[j], ctx.device, rb[j], re[j])) { return {}; } + } + for (int j = 0; j < 5; ++j) { + if (!writes[j]) { continue; } + if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } + for (int k = 0; k < 4; ++k) { if (wb[j] < re[k] && rb[k] < we[j]) { return {}; } } + for (int k = 0; k < j; ++k) { + if (writes[k] && wb[j] < we[k] && wb[k] < we[j] && + (wb[j] != wb[k] || we[j] != we[k] || writes[j]->type != writes[k]->type)) { return {}; } + } + } + } + return {mm, affine, count}; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr, const void * prepared_src1 = nullptr) { @@ -5901,6 +5954,16 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_cublas_affine(*cuda_ctx, cgraph, i); + if (match.count) { + const auto & affine = match.affine; + const ggml_cuda_affine_unary_ops ops = {affine.mul, affine.add, affine.unary, affine.post}; + ggml_cuda_mul_mat_cublas(*cuda_ctx, match.mm->src[0], match.mm->src[1], match.mm, prepared_src1, nullptr, nullptr, &ops); + return match.count - 1; + } + } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { const auto match = ggml_cuda_match_cublas_unary(*cuda_ctx, cgraph, i); if (match.count) { @@ -6614,6 +6677,13 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // add alloc deps for performance positive fusions. This may increase the overall compute buffer size. // TODO: consolidate fusion paths in graph_optimize and graph_compute for (int i = 0; i < cgraph->n_nodes; ++i) { + const auto cublas_affine = ggml_cuda_match_cublas_affine(*cuda_ctx, cgraph, i, false); + if (cublas_affine.count) { + ggml_tensor * terminal = cublas_affine.affine.post ? cublas_affine.affine.post : cublas_affine.affine.unary; + for (ggml_tensor * read : {cublas_affine.mm->src[0], cublas_affine.mm->src[1], cublas_affine.affine.mul->src[1], cublas_affine.affine.add->src[1]}) { + params->add_alloc_dep(params->user_data, read, terminal); + } + } const auto unary = ggml_cuda_match_cublas_unary(*cuda_ctx, cgraph, i, false); if (unary.count) { params->add_alloc_dep(params->user_data, unary.mm->src[0], unary.unary.last); diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 56affe6e614..9285e917c9c 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -247,7 +247,25 @@ static __device__ __forceinline__ float affine_unary_store(void * dst, ggml_type return value; } -static __global__ void affine_unary_kernel(affine_unary_args args) { +struct affine_unary_load { + __device__ __forceinline__ float operator()(const affine_unary_tensor & tensor, const uint32_t * pos, int) const { + return affine_unary_read(tensor, pos); + } +}; + +struct affine_unary_convert_load { + const void * data; + ggml_type type; + float * raw; + __device__ __forceinline__ float operator()(const affine_unary_tensor &, const uint32_t *, int i) const { + const float value = type == GGML_TYPE_F16 ? ggml_cuda_cast(((const half *) data)[i]) : ggml_cuda_cast(((const nv_bfloat16 *) data)[i]); + raw[i] = value; + return value; + } +}; + +template +static __device__ __forceinline__ void affine_unary_impl(affine_unary_args args, Load load) { ggml_cuda_pdl_lc(); const int i = blockDim.x*blockIdx.x + threadIdx.x; if (i >= args.count) { @@ -260,7 +278,7 @@ static __global__ void affine_unary_kernel(affine_unary_args args) { remaining /= args.ne[d]; } ggml_cuda_pdl_sync(); - float value = __fmul_rn(affine_unary_read(args.input[0], pos), affine_unary_read(args.input[1], pos)); + float value = __fmul_rn(load(args.input[0], pos, i), affine_unary_read(args.input[1], pos)); value = affine_unary_store(args.output[0], args.output_type[0], i, value); value = __fadd_rn(value, affine_unary_read(args.input[2], pos)); value = affine_unary_store(args.output[1], args.output_type[1], i, value); @@ -271,7 +289,15 @@ static __global__ void affine_unary_kernel(affine_unary_args args) { } } -void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { +static __global__ void affine_unary_kernel(affine_unary_args args) { + affine_unary_impl(args, affine_unary_load{}); +} + +static __global__ void affine_unary_convert_kernel(affine_unary_args args, const void * src, ggml_type type, float * raw) { + affine_unary_impl(args, affine_unary_convert_load{src, type, raw}); +} + +static affine_unary_args affine_unary_get_args(ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { affine_unary_args args{}; const ggml_tensor * input[] = { mul->src[0], mul->src[1], add->src[1] }; const ggml_tensor * output[] = { mul, add, unary, post }; @@ -296,11 +322,26 @@ void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mu memcpy(&args.scale, post->op_params, sizeof(float)); memcpy(&args.bias, (const char *) post->op_params + sizeof(float), sizeof(float)); } + return args; +} + +void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { + const auto args = affine_unary_get_args(mul, add, unary, post); const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); ggml_cuda_kernel_launch(affine_unary_kernel, launch, args); } +void ggml_cuda_op_affine_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_affine_unary_ops & ops) { + GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_BF16); + GGML_ASSERT(mm->type == GGML_TYPE_F32 && ggml_is_contiguous(mm) && ops.mul->src[0] == mm); + const auto args = affine_unary_get_args(ops.mul, ops.add, ops.unary, ops.post); + GGML_ASSERT(args.count == ggml_nelements(mm)); + const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(affine_unary_convert_kernel, launch, args, src, type, (float *) mm->data); +} + void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary(ctx, dst); } diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index fc0f124515b..fe6fc0a04c3 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -133,3 +133,12 @@ struct ggml_cuda_scaled_unary_args { }; void ggml_cuda_op_scaled_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_scaled_unary_args & args); + +struct ggml_cuda_affine_unary_ops { + ggml_tensor * mul; + ggml_tensor * add; + ggml_tensor * unary; + ggml_tensor * post; +}; + +void ggml_cuda_op_affine_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_affine_unary_ops & ops); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 05f7c6c2441..4f8b77124c9 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4154,19 +4154,20 @@ struct test_mul_mat_unary : public test_case { const bool before; const bool after; const bool batched; + const bool affine; std::vector stages; ggml_tensor * weights = nullptr; ggml_tensor * input = nullptr; - test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched) - : weight_type(weight_type), op(op), before(before), after(after), batched(batched) {} + test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched, bool affine = false) + : weight_type(weight_type), op(op), before(before), after(after), batched(batched), affine(affine) {} std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); return "MUL_MAT_UNARY"; } - std::string vars() override { return VARS_TO_STR5(weight_type, op, before, after, batched); } + std::string vars() override { return VARS_TO_STR6(weight_type, op, before, after, batched, affine); } bool run_whole_graph() override { return true; } std::vector fusion_test_nodes() override { return stages; } @@ -4175,6 +4176,14 @@ struct test_mul_mat_unary : public test_case { input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 127, 17, batched ? 2 : 1, 1); ggml_tensor * out = ggml_mul_mat(ctx, weights, input); stages = {out}; + if (affine) { + ggml_tensor * coefficient = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 129); + out = ggml_mul(ctx, out, coefficient); + stages.push_back(out); + out = ggml_add(ctx, out, bias); + stages.push_back(out); + } if (before) { out = ggml_scale_bias(ctx, out, 0.75f, 0.125f); stages.push_back(out); } out = ggml_unary(ctx, out, op); stages.push_back(out); @@ -9337,6 +9346,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_unary(type, op, true, true, false)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true)); } } From 33355933895daf4aafb12bc524fc94de7d95fee7 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 12:56:04 -0400 Subject: [PATCH 10/27] cuda : fuse tensor affine gates on cuBLAS result windows Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 17 +++++++++--- ggml/src/ggml-cuda/unary.cu | 48 ++++++++++++++++++++++++++++++++- tests/test-backend-ops.cpp | 16 ++++++++--- 3 files changed, 73 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 4fee77cd3ce..858daf2dc45 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5121,20 +5121,31 @@ struct ggml_cuda_cublas_affine_match { int count = 0; }; +static bool ggml_cuda_cublas_affine_input(const ggml_tensor * mm, const ggml_tensor * input) { + if (input == mm) { return true; } + if (!input || input->op != GGML_OP_VIEW || input->type != GGML_TYPE_F32 || input->view_src != mm || + input->view_offs % sizeof(float) || input->nb[0] != sizeof(float) || input->ne[0] <= 0 || input->ne[0] > mm->ne[0] || + input->view_offs/sizeof(float) > uint64_t(mm->ne[0] - input->ne[0])) { return false; } + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + if (input->ne[d] != mm->ne[d] || input->nb[d] != mm->nb[d]) { return false; } + } + return true; +} + static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { ggml_tensor * mm = graph->nodes[i]; if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || mm->src[1]->type != GGML_TYPE_F32 || - !ggml_is_contiguous(mm) || i + 1 >= graph->n_nodes || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } + !ggml_is_contiguous(mm) || ggml_nelements(mm) <= 0 || ggml_nelements(mm) > INT_MAX - CUDA_NEG_BLOCK_SIZE || i + 1 >= graph->n_nodes || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } int first = i + 1; while (first < graph->n_nodes && graph->nodes[first]->op == GGML_OP_VIEW) { - if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[first])) { return {}; } + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[first]) && !ggml_cuda_cublas_affine_input(mm, graph->nodes[first])) { return {}; } uintptr_t begin, end; if (allocated && !ggml_cuda_affine_unary_range(graph->nodes[first], ctx.device, begin, end)) { return {}; } ++first; } if (first >= graph->n_nodes) { return {}; } const auto affine = ggml_cuda_match_affine_unary(ctx, graph, first, allocated); - if (!affine.count || affine.mul->src[0] != mm || ggml_cuda_match_hc_affine_injection(ctx, graph, first, false).count) { return {}; } + if (!affine.count || !ggml_cuda_cublas_affine_input(mm, affine.mul->src[0]) || ggml_cuda_match_hc_affine_injection(ctx, graph, first, false).count) { return {}; } if (ggml_cuda_mul_mat_kernel(mm->src[0], mm->src[1], mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return {}; } const int cc = ggml_cuda_info().devices[ctx.device].cc; const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(cc, mm->src[0], mm->src[1], mm, false); diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 9285e917c9c..80357eaa781 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -264,6 +264,18 @@ struct affine_unary_convert_load { } }; +static __device__ __forceinline__ void affine_unary_apply(affine_unary_args args, float value, const uint32_t * pos, int i) { + value = __fmul_rn(value, affine_unary_read(args.input[1], pos)); + value = affine_unary_store(args.output[0], args.output_type[0], i, value); + value = __fadd_rn(value, affine_unary_read(args.input[2], pos)); + value = affine_unary_store(args.output[1], args.output_type[1], i, value); + value = args.op == GGML_UNARY_OP_SIGMOID ? op_sigmoid(value) : op_silu(value); + value = affine_unary_store(args.output[2], args.output_type[2], i, value); + if (args.output[3]) { + affine_unary_store(args.output[3], args.output_type[3], i, args.scale*value + args.bias); + } +} + template static __device__ __forceinline__ void affine_unary_impl(affine_unary_args args, Load load) { ggml_cuda_pdl_lc(); @@ -297,6 +309,25 @@ static __global__ void affine_unary_convert_kernel(affine_unary_args args, const affine_unary_impl(args, affine_unary_convert_load{src, type, raw}); } +static __global__ void affine_unary_window_convert_kernel(affine_unary_args args, const void * src, ggml_type type, float * raw, int count, int columns, int offset) { + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= count) { return; } + ggml_cuda_pdl_sync(); + const affine_unary_convert_load load{src, type, raw}; + const float value = load(args.input[0], nullptr, i); + const int column = i % columns; + if (column < offset || uint32_t(column - offset) >= args.ne[0]) { return; } + const int j = (i/columns)*args.ne[0] + column - offset; + uint32_t remaining = j; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + pos[d] = remaining % args.ne[d]; + remaining /= args.ne[d]; + } + affine_unary_apply(args, value, pos, j); +} + static affine_unary_args affine_unary_get_args(ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { affine_unary_args args{}; const ggml_tensor * input[] = { mul->src[0], mul->src[1], add->src[1] }; @@ -334,8 +365,23 @@ void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mu void ggml_cuda_op_affine_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_affine_unary_ops & ops) { GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_BF16); - GGML_ASSERT(mm->type == GGML_TYPE_F32 && ggml_is_contiguous(mm) && ops.mul->src[0] == mm); + GGML_ASSERT(mm->type == GGML_TYPE_F32 && ggml_is_contiguous(mm)); const auto args = affine_unary_get_args(ops.mul, ops.add, ops.unary, ops.post); + if (ops.mul->src[0] != mm) { + const ggml_tensor * input = ops.mul->src[0]; + GGML_ASSERT(input->op == GGML_OP_VIEW && input->view_src == mm && input->type == GGML_TYPE_F32 && + input->nb[0] == sizeof(float) && input->view_offs % sizeof(float) == 0 && input->ne[0] <= mm->ne[0] && + input->view_offs/sizeof(float) <= uint64_t(mm->ne[0] - input->ne[0])); + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + GGML_ASSERT(input->ne[d] == mm->ne[d] && input->nb[d] == mm->nb[d]); + } + GGML_ASSERT(ggml_nelements(mm) > 0 && ggml_nelements(mm) <= INT_MAX - CUDA_NEG_BLOCK_SIZE && args.count == ggml_nelements(input)); + const int count = ggml_nelements(mm); + const int blocks = (count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(affine_unary_window_convert_kernel, launch, args, src, type, (float *) mm->data, count, (int) mm->ne[0], (int) (input->view_offs/sizeof(float))); + return; + } GGML_ASSERT(args.count == ggml_nelements(mm)); const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4f8b77124c9..21af056785a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4155,19 +4155,20 @@ struct test_mul_mat_unary : public test_case { const bool after; const bool batched; const bool affine; + const bool window; std::vector stages; ggml_tensor * weights = nullptr; ggml_tensor * input = nullptr; - test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched, bool affine = false) - : weight_type(weight_type), op(op), before(before), after(after), batched(batched), affine(affine) {} + test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched, bool affine = false, bool window = false) + : weight_type(weight_type), op(op), before(before), after(after), batched(batched), affine(affine), window(window) {} std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); return "MUL_MAT_UNARY"; } - std::string vars() override { return VARS_TO_STR6(weight_type, op, before, after, batched, affine); } + std::string vars() override { return VARS_TO_STR7(weight_type, op, before, after, batched, affine, window); } bool run_whole_graph() override { return true; } std::vector fusion_test_nodes() override { return stages; } @@ -4176,9 +4177,14 @@ struct test_mul_mat_unary : public test_case { input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 127, 17, batched ? 2 : 1, 1); ggml_tensor * out = ggml_mul_mat(ctx, weights, input); stages = {out}; + if (window) { + GGML_ASSERT(affine); + out = ggml_view_4d(ctx, out, 71, out->ne[1], out->ne[2], out->ne[3], out->nb[1], out->nb[2], out->nb[3], 7*sizeof(float)); + stages.push_back(out); + } if (affine) { ggml_tensor * coefficient = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 129); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out->ne[0]); out = ggml_mul(ctx, out, coefficient); stages.push_back(out); out = ggml_add(ctx, out, bias); @@ -9348,6 +9354,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false, true)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false, true, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true, true)); } } From 5973fcae8c021702f2313722e32354abbe0e1321 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 14:29:46 -0400 Subject: [PATCH 11/27] cuda : fuse a second SCALE in tensor affine gates Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 57 +++++++++++++++++------------ ggml/src/ggml-cuda/unary.cu | 63 ++++++++++++++++++++++++++++++++- ggml/src/ggml-cuda/unary.cuh | 3 +- tests/test-backend-ops.cpp | 14 ++++++-- 4 files changed, 109 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 858daf2dc45..41ba5c811b6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4865,6 +4865,7 @@ struct ggml_cuda_affine_unary_match { ggml_tensor * add = nullptr; ggml_tensor * unary = nullptr; ggml_tensor * post = nullptr; + ggml_tensor * tail = nullptr; int count = 0; }; @@ -4893,7 +4894,7 @@ static bool ggml_cuda_affine_unary_prior_view(ggml_cgraph * graph, int start, co return false; } -static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { +static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true, bool with_tail = false) { ggml_cuda_affine_unary_match match; ggml_tensor * mul = graph->nodes[i]; if (mul->op != GGML_OP_MUL || !ggml_cuda_affine_unary_binary_type(mul)) { return match; } @@ -4922,9 +4923,12 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu const int p = next_compute(u + 1); ggml_tensor * post = p < graph->n_nodes && graph->nodes[p]->op == GGML_OP_SCALE && graph->nodes[p]->src[0] == unary ? graph->nodes[p] : nullptr; if (post && ((post->type != GGML_TYPE_F32 && post->type != GGML_TYPE_BF16) || post->type != unary->type)) { return match; } - const int indices[] = { i, a, u, p }; - const ggml_op ops[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY, GGML_OP_SCALE }; - const int n = post ? 4 : 3; + const int q = with_tail && post ? next_compute(p + 1) : graph->n_nodes; + ggml_tensor * tail = q < graph->n_nodes && graph->nodes[q]->op == GGML_OP_SCALE && graph->nodes[q]->src[0] == post ? graph->nodes[q] : nullptr; + if (tail && (tail->type != post->type || !ggml_is_contiguous(tail) || !ggml_are_same_shape(mul, tail))) { tail = nullptr; } + const int indices[] = { i, a, u, p, q }; + const ggml_op ops[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_UNARY, GGML_OP_SCALE, GGML_OP_SCALE }; + const int n = tail ? 5 : post ? 4 : 3; if (!ggml_can_fuse_subgraph_ext(graph, indices, n, ops, indices, n)) { return match; } int64_t elements = 1; for (int d = 0; d < GGML_MAX_DIMS; ++d) { @@ -4932,7 +4936,7 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu elements *= mul->ne[d]; } const ggml_tensor * reads[] = { mul->src[0], mul->src[1], add->src[1] }; - uintptr_t read_begin[3], read_end[3], begin[4], end[4]; + uintptr_t read_begin[3], read_end[3], begin[5], end[5]; for (int j = 0; j < 3; ++j) { const ggml_tensor * tensor = reads[j]; size_t bytes; @@ -4971,6 +4975,7 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu match.add = add; match.unary = unary; match.post = post; + match.tail = tail; match.count = indices[n - 1] - i + 1; return match; } @@ -5132,7 +5137,7 @@ static bool ggml_cuda_cublas_affine_input(const ggml_tensor * mm, const ggml_ten return true; } -static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { +static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true, bool with_tail = true) { ggml_tensor * mm = graph->nodes[i]; if (mm->op != GGML_OP_MUL_MAT || !mm->src[0] || !mm->src[1] || mm->type != GGML_TYPE_F32 || mm->src[1]->type != GGML_TYPE_F32 || !ggml_is_contiguous(mm) || ggml_nelements(mm) <= 0 || ggml_nelements(mm) > INT_MAX - CUDA_NEG_BLOCK_SIZE || i + 1 >= graph->n_nodes || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } @@ -5144,31 +5149,35 @@ static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_ ++first; } if (first >= graph->n_nodes) { return {}; } - const auto affine = ggml_cuda_match_affine_unary(ctx, graph, first, allocated); - if (!affine.count || !ggml_cuda_cublas_affine_input(mm, affine.mul->src[0]) || ggml_cuda_match_hc_affine_injection(ctx, graph, first, false).count) { return {}; } - if (ggml_cuda_mul_mat_kernel(mm->src[0], mm->src[1], mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return {}; } + auto affine = ggml_cuda_match_affine_unary(ctx, graph, first, allocated, with_tail); + if (!affine.count) { affine = ggml_cuda_match_affine_unary(ctx, graph, first, allocated); } + const auto reject = [&]() { + return with_tail && affine.tail ? ggml_cuda_match_cublas_affine(ctx, graph, i, allocated, false) : ggml_cuda_cublas_affine_match{}; + }; + if (!affine.count || !ggml_cuda_cublas_affine_input(mm, affine.mul->src[0]) || ggml_cuda_match_hc_affine_injection(ctx, graph, first, false).count) { return reject(); } + if (ggml_cuda_mul_mat_kernel(mm->src[0], mm->src[1], mm, ctx.device) != GGML_CUDA_MM_CUBLAS) { return reject(); } const int cc = ggml_cuda_info().devices[ctx.device].cc; const ggml_type type = ggml_cuda_mul_mat_cublas_compute_type(cc, mm->src[0], mm->src[1], mm, false); - if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || ggml_cuda_cublas_prefer_f32_output(type, cc)) { return {}; } + if ((type != GGML_TYPE_F16 && type != GGML_TYPE_BF16) || ggml_cuda_cublas_prefer_f32_output(type, cc)) { return reject(); } const int count = first - i + affine.count; std::vector ops; std::vector outputs; for (int j = i; j < i + count; ++j) { ops.push_back(graph->nodes[j]->op); outputs.push_back(j); } - if (!ggml_can_fuse_subgraph_ext(graph, outputs.data(), count, ops.data(), outputs.data(), outputs.size())) { return {}; } + if (!ggml_can_fuse_subgraph_ext(graph, outputs.data(), count, ops.data(), outputs.data(), outputs.size())) { return reject(); } if (allocated) { - uintptr_t rb[4], re[4], wb[5], we[5]; + uintptr_t rb[4], re[4], wb[6], we[6]; const ggml_tensor * reads[] = {mm->src[0], mm->src[1], affine.mul->src[1], affine.add->src[1]}; - const ggml_tensor * writes[] = {mm, affine.mul, affine.add, affine.unary, affine.post}; + const ggml_tensor * writes[] = {mm, affine.mul, affine.add, affine.unary, affine.post, affine.tail}; for (int j = 0; j < 4; ++j) { - if (j < 2 ? !ggml_cuda_hc_up_range(reads[j], ctx.device, rb[j], re[j]) : !ggml_cuda_affine_unary_range(reads[j], ctx.device, rb[j], re[j])) { return {}; } + if (j < 2 ? !ggml_cuda_hc_up_range(reads[j], ctx.device, rb[j], re[j]) : !ggml_cuda_affine_unary_range(reads[j], ctx.device, rb[j], re[j])) { return reject(); } } - for (int j = 0; j < 5; ++j) { + for (int j = 0; j < 6; ++j) { if (!writes[j]) { continue; } - if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } - for (int k = 0; k < 4; ++k) { if (wb[j] < re[k] && rb[k] < we[j]) { return {}; } } + if (!(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || !ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j])) { return reject(); } + for (int k = 0; k < 4; ++k) { if (wb[j] < re[k] && rb[k] < we[j]) { return reject(); } } for (int k = 0; k < j; ++k) { if (writes[k] && wb[j] < we[k] && wb[k] < we[j] && - (wb[j] != wb[k] || we[j] != we[k] || writes[j]->type != writes[k]->type)) { return {}; } + (wb[j] != wb[k] || we[j] != we[k] || writes[j]->type != writes[k]->type)) { return reject(); } } } } @@ -5969,7 +5978,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const auto match = ggml_cuda_match_cublas_affine(*cuda_ctx, cgraph, i); if (match.count) { const auto & affine = match.affine; - const ggml_cuda_affine_unary_ops ops = {affine.mul, affine.add, affine.unary, affine.post}; + const ggml_cuda_affine_unary_ops ops = {affine.mul, affine.add, affine.unary, affine.post, affine.tail}; ggml_cuda_mul_mat_cublas(*cuda_ctx, match.mm->src[0], match.mm->src[1], match.mm, prepared_src1, nullptr, nullptr, &ops); return match.count - 1; } @@ -6014,9 +6023,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { - const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); + auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, true, true); + if (!affine.count) { affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); } if (affine.count) { - ggml_cuda_op_affine_unary(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post); + ggml_cuda_op_affine_unary(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post, affine.tail); return affine.count - 1; } } @@ -6690,7 +6700,7 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph for (int i = 0; i < cgraph->n_nodes; ++i) { const auto cublas_affine = ggml_cuda_match_cublas_affine(*cuda_ctx, cgraph, i, false); if (cublas_affine.count) { - ggml_tensor * terminal = cublas_affine.affine.post ? cublas_affine.affine.post : cublas_affine.affine.unary; + ggml_tensor * terminal = cublas_affine.affine.tail ? cublas_affine.affine.tail : cublas_affine.affine.post ? cublas_affine.affine.post : cublas_affine.affine.unary; for (ggml_tensor * read : {cublas_affine.mm->src[0], cublas_affine.mm->src[1], cublas_affine.affine.mul->src[1], cublas_affine.affine.add->src[1]}) { params->add_alloc_dep(params->user_data, read, terminal); } @@ -6723,7 +6733,8 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph i += affine_injection.count - 1; continue; } - const auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); + auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false, true); + if (!affine.count) { affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, false); } if (affine.count) { ggml_tensor * last = cgraph->nodes[i + affine.count - 1]; params->add_alloc_dep(params->user_data, affine.mul->src[0], last); diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 80357eaa781..edb7d22736e 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -328,6 +328,38 @@ static __global__ void affine_unary_window_convert_kernel(affine_unary_args args affine_unary_apply(args, value, pos, j); } +static __global__ void affine_unary_tail_kernel(affine_unary_args args, void * tail, ggml_type tail_type, float scale, float bias, + const void * src, ggml_type src_type, float * raw, int count, int columns, int offset) { + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= count) { return; } + ggml_cuda_pdl_sync(); + float value = 0.0f; + int j = i; + if (raw) { + const affine_unary_convert_load load{src, src_type, raw}; + value = load(args.input[0], nullptr, i); + const int column = i % columns; + if (column < offset || uint32_t(column - offset) >= args.ne[0]) { return; } + j = (i/columns)*args.ne[0] + column - offset; + } + uint32_t remaining = j; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + pos[d] = remaining % args.ne[d]; + remaining /= args.ne[d]; + } + if (!raw) { value = affine_unary_read(args.input[0], pos); } + value = __fmul_rn(value, affine_unary_read(args.input[1], pos)); + value = affine_unary_store(args.output[0], args.output_type[0], j, value); + value = __fadd_rn(value, affine_unary_read(args.input[2], pos)); + value = affine_unary_store(args.output[1], args.output_type[1], j, value); + value = args.op == GGML_UNARY_OP_SIGMOID ? op_sigmoid(value) : op_silu(value); + value = affine_unary_store(args.output[2], args.output_type[2], j, value); + value = affine_unary_store(args.output[3], args.output_type[3], j, args.scale*value + args.bias); + affine_unary_store(tail, tail_type, j, scale*value + bias); +} + static affine_unary_args affine_unary_get_args(ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { affine_unary_args args{}; const ggml_tensor * input[] = { mul->src[0], mul->src[1], add->src[1] }; @@ -356,8 +388,25 @@ static affine_unary_args affine_unary_get_args(ggml_tensor * mul, ggml_tensor * return args; } -void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post) { +static void affine_unary_tail_cuda(ggml_backend_cuda_context & ctx, affine_unary_args args, ggml_tensor * tail, + const void * src = nullptr, ggml_type src_type = GGML_TYPE_F32, ggml_tensor * mm = nullptr, int offset = 0) { + GGML_ASSERT(args.output[3] && tail && (tail->type == GGML_TYPE_F32 || tail->type == GGML_TYPE_BF16)); + float scale, bias; + memcpy(&scale, tail->op_params, sizeof(float)); + memcpy(&bias, (const char *) tail->op_params + sizeof(float), sizeof(float)); + const int count = mm ? ggml_nelements(mm) : args.count; + const int blocks = (count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(affine_unary_tail_kernel, launch, args, tail->data, tail->type, scale, bias, + src, src_type, mm ? (float *) mm->data : nullptr, count, mm ? (int) mm->ne[0] : 0, offset); +} + +void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post, ggml_tensor * tail) { const auto args = affine_unary_get_args(mul, add, unary, post); + if (tail) { + affine_unary_tail_cuda(ctx, args, tail); + return; + } const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); ggml_cuda_kernel_launch(affine_unary_kernel, launch, args); @@ -367,6 +416,18 @@ void ggml_cuda_op_affine_unary_convert(ggml_backend_cuda_context & ctx, ggml_typ GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_BF16); GGML_ASSERT(mm->type == GGML_TYPE_F32 && ggml_is_contiguous(mm)); const auto args = affine_unary_get_args(ops.mul, ops.add, ops.unary, ops.post); + if (ops.tail) { + const ggml_tensor * input = ops.mul->src[0]; + GGML_ASSERT(input == mm || (input->op == GGML_OP_VIEW && input->view_src == mm && input->type == GGML_TYPE_F32 && + input->nb[0] == sizeof(float) && input->view_offs % sizeof(float) == 0 && input->ne[0] > 0 && input->ne[0] <= mm->ne[0] && + input->view_offs/sizeof(float) <= uint64_t(mm->ne[0] - input->ne[0]))); + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + GGML_ASSERT(input->ne[d] == mm->ne[d] && input->nb[d] == mm->nb[d]); + } + GGML_ASSERT(ggml_nelements(mm) > 0 && ggml_nelements(mm) <= INT_MAX - CUDA_NEG_BLOCK_SIZE && args.count == ggml_nelements(input)); + affine_unary_tail_cuda(ctx, args, ops.tail, src, type, mm, input == mm ? 0 : (int) (input->view_offs/sizeof(float))); + return; + } if (ops.mul->src[0] != mm) { const ggml_tensor * input = ops.mul->src[0]; GGML_ASSERT(input->op == GGML_OP_VIEW && input->view_src == mm && input->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index fe6fc0a04c3..0148664e740 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -93,7 +93,7 @@ void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); -void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post); +void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post, ggml_tensor * tail = nullptr); void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); @@ -139,6 +139,7 @@ struct ggml_cuda_affine_unary_ops { ggml_tensor * add; ggml_tensor * unary; ggml_tensor * post; + ggml_tensor * tail = nullptr; }; void ggml_cuda_op_affine_unary_convert(ggml_backend_cuda_context & ctx, ggml_type type, const void * src, ggml_tensor * mm, const ggml_cuda_affine_unary_ops & ops); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 21af056785a..fd9ef83d32f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4156,19 +4156,20 @@ struct test_mul_mat_unary : public test_case { const bool batched; const bool affine; const bool window; + const bool second_post; std::vector stages; ggml_tensor * weights = nullptr; ggml_tensor * input = nullptr; - test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched, bool affine = false, bool window = false) - : weight_type(weight_type), op(op), before(before), after(after), batched(batched), affine(affine), window(window) {} + test_mul_mat_unary(ggml_type weight_type, ggml_unary_op op, bool before, bool after, bool batched, bool affine = false, bool window = false, bool second_post = false) + : weight_type(weight_type), op(op), before(before), after(after), batched(batched), affine(affine), window(window), second_post(second_post) {} std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); return "MUL_MAT_UNARY"; } - std::string vars() override { return VARS_TO_STR7(weight_type, op, before, after, batched, affine, window); } + std::string vars() override { return VARS_TO_STR8(weight_type, op, before, after, batched, affine, window, second_post); } bool run_whole_graph() override { return true; } std::vector fusion_test_nodes() override { return stages; } @@ -4194,6 +4195,11 @@ struct test_mul_mat_unary : public test_case { out = ggml_unary(ctx, out, op); stages.push_back(out); if (after) { out = ggml_scale_bias(ctx, out, 1.5f, -0.025f); stages.push_back(out); } + if (second_post) { + GGML_ASSERT(affine && after); + out = ggml_scale_bias(ctx, out, 0.75f, 0.015625f); + stages.push_back(out); + } for (ggml_tensor * stage : stages) { ggml_set_output(stage); } ggml_set_name(out, "out"); return out; @@ -9356,6 +9362,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, false, false, true, true)); test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, false, true, false, true)); + test_cases.emplace_back(new test_mul_mat_unary(type, op, false, true, true, true, true, true)); } } From ce6ddc95589080cf0a0cf0c0fa42652fa962f940 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 16:18:50 -0400 Subject: [PATCH 12/27] cuda : fuse ordinary residual MUL and ADD Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 70 +++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/unary.cu | 65 ++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/unary.cuh | 2 + tests/test-backend-ops.cpp | 36 +++++++++++++++-- 4 files changed, 170 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 41ba5c811b6..bc5cdc89ce6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4980,6 +4980,62 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu return match; } +static ggml_cuda_affine_unary_match ggml_cuda_match_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * mul = graph->nodes[i]; + if (mul->op != GGML_OP_MUL || !ggml_cuda_affine_unary_binary_type(mul)) { return {}; } + int a = i + 1; + while (a < graph->n_nodes && graph->nodes[a]->op == GGML_OP_VIEW) { + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[a]) || + (allocated && !ggml_cuda_affine_unary_range(graph->nodes[a], ctx.device, begin, end))) { return {}; } + ++a; + } + if (a >= graph->n_nodes || graph->nodes[a]->op != GGML_OP_ADD || !ggml_cuda_affine_unary_binary_type(graph->nodes[a])) { return {}; } + ggml_tensor * add = graph->nodes[a]; + if ((add->src[0] == mul) == (add->src[1] == mul)) { return {}; } + if (a + 1 < graph->n_nodes && ggml_can_fuse(graph, a, {GGML_OP_ADD, GGML_OP_ADD}) && + graph->nodes[a + 1]->src[0] == add && ggml_are_same_layout(add->src[1], graph->nodes[a + 1]->src[1])) { return {}; } + const int indices[] = {i, a}; + const ggml_op ops[] = {GGML_OP_MUL, GGML_OP_ADD}; + if (!ggml_can_fuse_subgraph_ext(graph, indices, 2, ops, indices, 2)) { return {}; } + int64_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (mul->ne[d] <= 0 || mul->ne[d] > (INT_MAX - CUDA_NEG_BLOCK_SIZE)/elements) { return {}; } + elements *= mul->ne[d]; + } + const ggml_tensor * reads[] = {mul->src[0], mul->src[1], add->src[add->src[0] == mul ? 1 : 0]}; + uintptr_t rb[3], re[3], wb[2], we[2]; + for (int j = 0; j < 3; ++j) { + const ggml_tensor * tensor = reads[j]; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(tensor, bytes) || (allocated && + (!ggml_cuda_affine_unary_range(tensor, ctx.device, rb[j], re[j]) || rb[j] % ggml_type_size(tensor->type)))) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (tensor->ne[d] > INT_MAX || tensor->nb[d] % ggml_type_size(tensor->type) || + tensor->nb[d]/ggml_type_size(tensor->type) > UINT32_MAX || mul->ne[d] % tensor->ne[d]) { return {}; } + } + if (tensor == mul || tensor == add || tensor->view_src == mul || tensor->view_src == add) { return {}; } + } + if (!ggml_are_same_shape(reads[0], mul)) { return {}; } + for (int j = 0; j < 2; ++j) { + const ggml_tensor * tensor = graph->nodes[indices[j]]; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(tensor, bytes) || !ggml_is_contiguous(tensor) || !ggml_are_same_shape(mul, tensor) || (allocated && + (!ggml_cuda_affine_unary_range(tensor, ctx.device, wb[j], we[j]) || wb[j] % ggml_type_size(tensor->type)))) { return {}; } + if (!allocated) { continue; } + for (int k = 0; k < 3; ++k) { + if (wb[j] < re[k] && rb[k] < we[j] && (k != 0 || wb[j] != rb[k] || we[j] != re[k] || + tensor->type != reads[0]->type || !ggml_is_contiguous(reads[0]))) { return {}; } + } + if (j && wb[j] < we[0] && wb[0] < we[j] && (wb[j] != wb[0] || we[j] != we[0] || tensor->type != mul->type)) { return {}; } + } + ggml_cuda_affine_unary_match match; + match.mul = mul; + match.add = add; + match.count = a - i + 1; + return match; +} + struct ggml_cuda_hc_affine_injection_match { ggml_cuda_affine_unary_match affine; ggml_tensor * post = nullptr; @@ -6029,6 +6085,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_cuda_op_affine_unary(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post, affine.tail); return affine.count - 1; } + const auto binary = ggml_cuda_match_mul_add(*cuda_ctx, cgraph, i); + if (binary.count) { + ggml_cuda_op_mul_add(*cuda_ctx, binary.mul, binary.add); + return binary.count - 1; + } } @@ -6744,6 +6805,15 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph continue; } + const auto binary = ggml_cuda_match_mul_add(*cuda_ctx, cgraph, i, false); + if (binary.count) { + for (ggml_tensor * read : {binary.mul->src[0], binary.mul->src[1], binary.add->src[binary.add->src[0] == binary.mul ? 1 : 0]}) { + params->add_alloc_dep(params->user_data, read, binary.add); + } + i += binary.count - 1; + continue; + } + const auto injection = ggml_cuda_match_hc_injection(*cuda_ctx, cgraph, i, false); if (injection.count) { ggml_tensor * terminal = cgraph->nodes[i + injection.count - 1]; diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index edb7d22736e..2790bd46332 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -388,6 +388,71 @@ static affine_unary_args affine_unary_get_args(ggml_tensor * mul, ggml_tensor * return args; } +struct mul_add_args { + affine_unary_args affine; + uint3 ne[GGML_MAX_DIMS]; + uint3 repeat[3][GGML_MAX_DIMS]; + uint32_t repeat_mask[3]; +}; + +static __device__ __forceinline__ float mul_add_read(const affine_unary_tensor & tensor, const uint3 * repeat, const uint32_t * pos, uint32_t repeat_mask) { + size_t offset = 0; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + const uint32_t index = repeat_mask & (1u << d) ? fastmodulo(pos[d], repeat[d]) : pos[d]; + offset += size_t(index)*tensor.nb[d]; + } + const char * ptr = (const char *) tensor.data + offset; + if (tensor.type == GGML_TYPE_F16) { return ggml_cuda_cast(*(const half *) ptr); } + if (tensor.type == GGML_TYPE_BF16) { return ggml_cuda_cast(*(const nv_bfloat16 *) ptr); } + return *(const float *) ptr; +} + +static __global__ void mul_add_kernel(mul_add_args binary, bool product_first) { + const affine_unary_args & args = binary.affine; + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= args.count) { return; } + uint32_t remaining = i; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + const uint2 part = fast_div_modulo(remaining, binary.ne[d]); + remaining = part.x; + pos[d] = part.y; + } + ggml_cuda_pdl_sync(); + float value = __fmul_rn(mul_add_read(args.input[0], binary.repeat[0], pos, 0), mul_add_read(args.input[1], binary.repeat[1], pos, binary.repeat_mask[1])); + value = affine_unary_store(args.output[0], args.output_type[0], i, value); + const float base = mul_add_read(args.input[2], binary.repeat[2], pos, binary.repeat_mask[2]); + value = product_first ? __fadd_rn(value, base) : __fadd_rn(base, value); + affine_unary_store(args.output[1], args.output_type[1], i, value); +} + +void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add) { + mul_add_args binary{}; + affine_unary_args & args = binary.affine; + const bool product_first = add->src[0] == mul; + const ggml_tensor * reads[] = {mul->src[0], mul->src[1], add->src[product_first ? 1 : 0]}; + for (int j = 0; j < 3; ++j) { + args.input[j].data = reads[j]->data; + args.input[j].type = reads[j]->type; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + args.input[j].ne[d] = reads[j]->ne[d]; + args.input[j].nb[d] = reads[j]->ne[d] == 1 ? 0 : reads[j]->nb[d]; + binary.repeat[j][d] = init_fastdiv_values(reads[j]->ne[d]); + if (reads[j]->ne[d] != 1 && reads[j]->ne[d] != mul->ne[d]) { binary.repeat_mask[j] |= 1u << d; } + } + } + args.output[0] = mul->data; + args.output[1] = add->data; + args.output_type[0] = mul->type; + args.output_type[1] = add->type; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { binary.ne[d] = init_fastdiv_values(mul->ne[d]); } + args.count = ggml_nelements(mul); + const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(mul_add_kernel, launch, binary, product_first); +} + static void affine_unary_tail_cuda(ggml_backend_cuda_context & ctx, affine_unary_args args, ggml_tensor * tail, const void * src = nullptr, ggml_type src_type = GGML_TYPE_F32, ggml_tensor * mm = nullptr, int offset = 0) { GGML_ASSERT(args.output[3] && tail && (tail->type == GGML_TYPE_F32 || tail->type == GGML_TYPE_BF16)); diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 0148664e740..4b8275a64e0 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -93,6 +93,8 @@ void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); +void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add); + void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post, ggml_tensor * tail = nullptr); void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index fd9ef83d32f..ef3e6f8d60d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3296,14 +3296,23 @@ struct test_bin_bcast : public test_case { bool perm1; // permute src1? bool src_overlap; // src0 and src1 are overlapping views of the same buffer const ggml_type type_b; // type of src1, GGML_TYPE_COUNT -> same as src0 + const bool add_after; + const bool swap_add; + std::vector stages; - bool run_whole_graph() override { return nf > 1; } + bool run_whole_graph() override { return nf > 1 || add_after; } + std::vector fusion_test_nodes() override { return stages; } + + std::string op_desc(ggml_tensor * t) override { + return add_after ? "MUL_ADD" : test_case::op_desc(t); + } std::string vars() override { std::string s = VARS_TO_STR6(type, ne, nr, nf, perm1, src_overlap); if (type_b != GGML_TYPE_COUNT) { s += "," + VAR_TO_STR(type_b); } + if (add_after) { s += "," + VARS_TO_STR2(add_after, swap_add); } return s; } @@ -3316,11 +3325,12 @@ struct test_bin_bcast : public test_case { std::array nr = {1, 2, 1, 1}, int nf = 1, bool perm1 = false, bool src_overlap = false, - ggml_type type_b = GGML_TYPE_COUNT) - : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b) {} + ggml_type type_b = GGML_TYPE_COUNT, bool add_after = false, bool swap_add = false) + : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b), add_after(add_after), swap_add(swap_add) {} ggml_tensor * build_graph(ggml_context * ctx) override { GGML_ASSERT(nf <= 16); + stages.clear(); ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]); ggml_set_name(a, "a"); @@ -3361,6 +3371,15 @@ struct test_bin_bcast : public test_case { out = op(ctx, out, b[i]); } + if (add_after) { + GGML_ASSERT(op == ggml_mul && nf == 1); + ggml_set_output(out); + stages.push_back(out); + ggml_tensor * residual = ggml_new_tensor(ctx, type, 4, out->ne); + ggml_set_name(residual, "residual"); + out = swap_add ? ggml_add(ctx, residual, out) : ggml_add(ctx, out, residual); + stages.push_back(out); + } ggml_set_name(out, "out"); return out; @@ -10007,6 +10026,17 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (bool permute : {false, true}) { + for (bool swap : {false, true}) { + test_cases.emplace_back(new test_bin_bcast(ggml_mul, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, permute, false, GGML_TYPE_COUNT, true, swap)); + if (type != GGML_TYPE_F32) { + test_cases.emplace_back(new test_bin_bcast(ggml_mul, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, permute, false, GGML_TYPE_F32, true, swap)); + } + } + } + } + auto add_test_bin_bcast = [&](ggml_type type, std::array ne, std::array nr, bool perm1 = false, bool src_overlap = false) { for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) { test_cases.emplace_back(new test_bin_bcast(op, type, ne, nr, 1, perm1, src_overlap)); From 0dead34ccc8c639a0df6d3ca91e19edbbc0d38e1 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 17:25:33 -0400 Subject: [PATCH 13/27] cuda : fuse repeated residual MUL and ADD Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 78 ++++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/unary.cu | 47 ++++++++++++++++++-- ggml/src/ggml-cuda/unary.cuh | 1 + tests/test-backend-ops.cpp | 37 +++++++++++++--- 4 files changed, 150 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index bc5cdc89ce6..a96cf36e2de 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4879,10 +4879,14 @@ static bool ggml_cuda_affine_unary_binary_type(const ggml_tensor * node) { (a == GGML_TYPE_BF16 && (b == GGML_TYPE_F32 || b == GGML_TYPE_BF16) && node->type == GGML_TYPE_BF16); } -static bool ggml_cuda_affine_unary_prior_view(ggml_cgraph * graph, int start, const ggml_tensor * tensor) { +static bool ggml_cuda_affine_unary_view_op(ggml_op op) { + return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; +} + +static bool ggml_cuda_affine_unary_prior_view(ggml_cgraph * graph, int start, const ggml_tensor * tensor, bool metadata = false) { const ggml_tensor * root = tensor->view_src; size_t bytes, root_bytes; - if (tensor->op != GGML_OP_VIEW || !root || root->view_src || + if ((tensor->op != GGML_OP_VIEW && (!metadata || !ggml_cuda_affine_unary_view_op(tensor->op))) || !root || root->view_src || !ggml_cuda_affine_unary_bytes(tensor, bytes) || !ggml_cuda_affine_unary_bytes(root, root_bytes) || tensor->view_offs > root_bytes || bytes > root_bytes - tensor->view_offs) { return false; @@ -4980,13 +4984,13 @@ static ggml_cuda_affine_unary_match ggml_cuda_match_affine_unary(ggml_backend_cu return match; } -static ggml_cuda_affine_unary_match ggml_cuda_match_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { +static ggml_cuda_affine_unary_match ggml_cuda_match_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true, bool metadata = false) { ggml_tensor * mul = graph->nodes[i]; if (mul->op != GGML_OP_MUL || !ggml_cuda_affine_unary_binary_type(mul)) { return {}; } int a = i + 1; - while (a < graph->n_nodes && graph->nodes[a]->op == GGML_OP_VIEW) { + while (a < graph->n_nodes && (metadata ? ggml_cuda_affine_unary_view_op(graph->nodes[a]->op) : graph->nodes[a]->op == GGML_OP_VIEW)) { uintptr_t begin, end; - if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[a]) || + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[a], metadata) || (allocated && !ggml_cuda_affine_unary_range(graph->nodes[a], ctx.device, begin, end))) { return {}; } ++a; } @@ -5240,6 +5244,57 @@ static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_ return {mm, affine, count}; } +struct ggml_cuda_repeat_mul_add_match { + ggml_tensor * repeat = nullptr; + ggml_cuda_affine_unary_match binary; + int count = 0; +}; + +static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * repeat = graph->nodes[i]; + if (repeat->op != GGML_OP_REPEAT || !repeat->src[0] || repeat->type != repeat->src[0]->type || + (repeat->type != GGML_TYPE_F32 && repeat->type != GGML_TYPE_F16)) { return {}; } + int m = i + 1; + while (m < graph->n_nodes && ggml_cuda_affine_unary_view_op(graph->nodes[m]->op)) { + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[m], true) || + (allocated && !ggml_cuda_affine_unary_range(graph->nodes[m], ctx.device, begin, end))) { return {}; } + ++m; + } + if (m >= graph->n_nodes || (graph->nodes[m]->src[0] == repeat) == (graph->nodes[m]->src[1] == repeat)) { return {}; } + if (ggml_cuda_match_hc_affine_injection(ctx, graph, m, false).count || ggml_cuda_match_affine_unary(ctx, graph, m, false).count) { return {}; } + const auto binary = ggml_cuda_match_mul_add(ctx, graph, m, allocated, true); + if (!binary.count || !ggml_are_same_shape(repeat, binary.mul) || !ggml_is_contiguous(repeat)) { return {}; } + const int indices[] = {i, m, m + binary.count - 1}; + const ggml_op ops[] = {GGML_OP_REPEAT, GGML_OP_MUL, GGML_OP_ADD}; + if (!ggml_can_fuse_subgraph_ext(graph, indices, 3, ops, indices, 3)) { return {}; } + const ggml_tensor * reads[] = {repeat->src[0], binary.mul->src[binary.mul->src[0] == repeat ? 1 : 0], binary.add->src[binary.add->src[0] == binary.mul ? 1 : 0]}; + const ggml_tensor * writes[] = {repeat, binary.mul, binary.add}; + uintptr_t rb[3], re[3], wb[3], we[3]; + for (int j = 0; j < 3; ++j) { + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(reads[j], bytes) || (allocated && + (!ggml_cuda_affine_unary_range(reads[j], ctx.device, rb[j], re[j]) || rb[j] % ggml_type_size(reads[j]->type)))) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (reads[j]->ne[d] > INT_MAX || reads[j]->nb[d] % ggml_type_size(reads[j]->type) || + reads[j]->nb[d]/ggml_type_size(reads[j]->type) > UINT32_MAX || binary.mul->ne[d] % reads[j]->ne[d]) { return {}; } + } + for (const ggml_tensor * write : writes) { + if (reads[j] == write || reads[j]->view_src == write) { return {}; } + } + } + if (allocated) { + for (int j = 0; j < 3; ++j) { + if (!ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j]) || wb[j] % ggml_type_size(writes[j]->type)) { return {}; } + for (int k = 0; k < 3; ++k) { if (wb[j] < re[k] && rb[k] < we[j]) { return {}; } } + for (int k = 0; k < j; ++k) { + if (wb[j] < we[k] && wb[k] < we[j] && (wb[j] != wb[k] || we[j] != we[k] || writes[j]->type != writes[k]->type)) { return {}; } + } + } + } + return {repeat, binary, indices[2] - i + 1}; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr, const void * prepared_src1 = nullptr) { @@ -6079,6 +6134,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto repeated = ggml_cuda_match_repeat_mul_add(*cuda_ctx, cgraph, i); + if (repeated.count) { + ggml_cuda_op_repeat_mul_add(*cuda_ctx, repeated.repeat, repeated.binary.mul, repeated.binary.add); + return repeated.count - 1; + } auto affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i, true, true); if (!affine.count) { affine = ggml_cuda_match_affine_unary(*cuda_ctx, cgraph, i); } if (affine.count) { @@ -6778,6 +6838,14 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto repeated = ggml_cuda_match_repeat_mul_add(*cuda_ctx, cgraph, i, false); + if (repeated.count) { + for (ggml_tensor * read : {repeated.repeat->src[0], repeated.repeat, repeated.binary.mul->src[0], repeated.binary.mul->src[1], repeated.binary.add->src[repeated.binary.add->src[0] == repeated.binary.mul ? 1 : 0]}) { + params->add_alloc_dep(params->user_data, read, repeated.binary.add); + } + i += repeated.count - 1; + continue; + } const auto affine_injection = ggml_cuda_match_hc_affine_injection(*cuda_ctx, cgraph, i, false); if (affine_injection.count) { ggml_tensor * terminal = cgraph->nodes[i + affine_injection.count - 1]; diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 2790bd46332..fc02803fd22 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -427,11 +427,36 @@ static __global__ void mul_add_kernel(mul_add_args binary, bool product_first) { affine_unary_store(args.output[1], args.output_type[1], i, value); } -void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add) { +static __global__ void repeat_mul_add_kernel(mul_add_args binary, int repeated_input, bool product_first) { + const affine_unary_args & args = binary.affine; + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= args.count) { return; } + uint32_t remaining = i; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + const uint2 part = fast_div_modulo(remaining, binary.ne[d]); + remaining = part.x; + pos[d] = part.y; + } + ggml_cuda_pdl_sync(); + float lhs = mul_add_read(args.input[0], binary.repeat[0], pos, binary.repeat_mask[0]); + float rhs = mul_add_read(args.input[1], binary.repeat[1], pos, binary.repeat_mask[1]); + if (repeated_input == 0) { lhs = affine_unary_store(args.output[2], args.output_type[2], i, lhs); } + else { rhs = affine_unary_store(args.output[2], args.output_type[2], i, rhs); } + float value = __fmul_rn(lhs, rhs); + value = affine_unary_store(args.output[0], args.output_type[0], i, value); + const float base = mul_add_read(args.input[2], binary.repeat[2], pos, binary.repeat_mask[2]); + value = product_first ? __fadd_rn(value, base) : __fadd_rn(base, value); + affine_unary_store(args.output[1], args.output_type[1], i, value); +} + +static mul_add_args mul_add_get_args(ggml_tensor * mul, ggml_tensor * add, ggml_tensor * repeat = nullptr) { mul_add_args binary{}; affine_unary_args & args = binary.affine; const bool product_first = add->src[0] == mul; const ggml_tensor * reads[] = {mul->src[0], mul->src[1], add->src[product_first ? 1 : 0]}; + if (repeat) { reads[mul->src[0] == repeat ? 0 : 1] = repeat->src[0]; } for (int j = 0; j < 3; ++j) { args.input[j].data = reads[j]->data; args.input[j].type = reads[j]->type; @@ -448,9 +473,25 @@ void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, gg args.output_type[1] = add->type; for (int d = 0; d < GGML_MAX_DIMS; ++d) { binary.ne[d] = init_fastdiv_values(mul->ne[d]); } args.count = ggml_nelements(mul); - const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + if (repeat) { + args.output[2] = repeat->data; + args.output_type[2] = repeat->type; + } + return binary; +} + +void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add) { + const auto binary = mul_add_get_args(mul, add); + const int blocks = (binary.affine.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(mul_add_kernel, launch, binary, add->src[0] == mul); +} + +void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * repeat, ggml_tensor * mul, ggml_tensor * add) { + const auto binary = mul_add_get_args(mul, add, repeat); + const int blocks = (binary.affine.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); - ggml_cuda_kernel_launch(mul_add_kernel, launch, binary, product_first); + ggml_cuda_kernel_launch(repeat_mul_add_kernel, launch, binary, mul->src[0] == repeat ? 0 : 1, add->src[0] == mul); } static void affine_unary_tail_cuda(ggml_backend_cuda_context & ctx, affine_unary_args args, ggml_tensor * tail, diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 4b8275a64e0..909ac2a47b3 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -94,6 +94,7 @@ void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add); +void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * repeat, ggml_tensor * mul, ggml_tensor * add); void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post, ggml_tensor * tail = nullptr); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ef3e6f8d60d..09fd82c804a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3298,13 +3298,14 @@ struct test_bin_bcast : public test_case { const ggml_type type_b; // type of src1, GGML_TYPE_COUNT -> same as src0 const bool add_after; const bool swap_add; + const int repeat_input; std::vector stages; bool run_whole_graph() override { return nf > 1 || add_after; } std::vector fusion_test_nodes() override { return stages; } std::string op_desc(ggml_tensor * t) override { - return add_after ? "MUL_ADD" : test_case::op_desc(t); + return repeat_input >= 0 ? "REPEAT_MUL_ADD" : add_after ? "MUL_ADD" : test_case::op_desc(t); } std::string vars() override { @@ -3313,6 +3314,7 @@ struct test_bin_bcast : public test_case { s += "," + VAR_TO_STR(type_b); } if (add_after) { s += "," + VARS_TO_STR2(add_after, swap_add); } + if (repeat_input >= 0) { s += "," + VAR_TO_STR(repeat_input); } return s; } @@ -3325,14 +3327,15 @@ struct test_bin_bcast : public test_case { std::array nr = {1, 2, 1, 1}, int nf = 1, bool perm1 = false, bool src_overlap = false, - ggml_type type_b = GGML_TYPE_COUNT, bool add_after = false, bool swap_add = false) - : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b), add_after(add_after), swap_add(swap_add) {} + ggml_type type_b = GGML_TYPE_COUNT, bool add_after = false, bool swap_add = false, int repeat_input = -1) + : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b), add_after(add_after), swap_add(swap_add), repeat_input(repeat_input) {} ggml_tensor * build_graph(ggml_context * ctx) override { GGML_ASSERT(nf <= 16); stages.clear(); - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]); + std::array expanded = {ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]}; + ggml_tensor * a = ggml_new_tensor(ctx, type, 4, repeat_input == 0 ? ne.data() : expanded.data()); ggml_set_name(a, "a"); const ggml_type tb = type_b == GGML_TYPE_COUNT ? type : type_b; @@ -3340,7 +3343,7 @@ struct test_bin_bcast : public test_case { ggml_tensor * b[16]; for (int i = 0; i < nf; ++i) { if (perm1) { - const int p[4] = { 1, 2, 0, 3 }; // hardcoded for now + const int p[4] = { repeat_input == 1 ? 0 : 1, 2, repeat_input == 1 ? 1 : 0, 3 }; // hardcoded for now b[i] = ggml_new_tensor_4d(ctx, tb, ne[p[0]], ne[p[1]], ne[p[2]], ne[p[3]]); b[i] = ggml_permute(ctx, b[i], p[0], p[1], p[2], p[3]); @@ -3352,6 +3355,15 @@ struct test_bin_bcast : public test_case { ggml_set_name(b[i], (std::string("b") + std::to_string(i)).c_str()); } + if (repeat_input >= 0) { + GGML_ASSERT(op == ggml_mul && nf == 1 && add_after && !src_overlap && repeat_input <= 1); + ggml_tensor * repeated = ggml_repeat_4d(ctx, repeat_input == 0 ? a : b[0], expanded[0], expanded[1], expanded[2], expanded[3]); + ggml_set_output(repeated); + stages.push_back(repeated); + if (repeat_input == 0) a = repeated; + else b[0] = repeated; + } + // The backward pass supports broadcasting only for GGML_ADD: const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; if (grad_supported) { @@ -10037,6 +10049,21 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (bool permute : {false, true}) { + for (bool swap : {false, true}) { + for (int repeated : {0, 1}) { + if (type != GGML_TYPE_BF16) { + test_cases.emplace_back(new test_bin_bcast(ggml_mul, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, permute, false, GGML_TYPE_COUNT, true, swap, repeated)); + } + if (type == GGML_TYPE_F16 || (type == GGML_TYPE_BF16 && repeated == 1)) { + test_cases.emplace_back(new test_bin_bcast(ggml_mul, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, permute, false, GGML_TYPE_F32, true, swap, repeated)); + } + } + } + } + } + auto add_test_bin_bcast = [&](ggml_type type, std::array ne, std::array nr, bool perm1 = false, bool src_overlap = false) { for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) { test_cases.emplace_back(new test_bin_bcast(op, type, ne, nr, 1, perm1, src_overlap)); From 24e00b647d842a1ae6a0248c447de36087b252e2 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 18:18:07 -0400 Subject: [PATCH 14/27] cuda : fuse ordered stream MUL and ADD Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 110 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/unary.cu | 57 +++++++++++++++++ ggml/src/ggml-cuda/unary.cuh | 3 + tests/test-backend-ops.cpp | 41 ++++++++++-- 4 files changed, 204 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a96cf36e2de..6757075931e 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5295,6 +5295,99 @@ static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backen return {repeat, binary, indices[2] - i + 1}; } +struct ggml_cuda_ordered_mul_add_match { + ggml_tensor * products[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; + ggml_tensor * sums[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; + int streams = 0; + int count = 0; +}; + +static ggml_cuda_ordered_mul_add_match ggml_cuda_match_ordered_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * first = graph->nodes[i]; + if (first->op != GGML_OP_MUL || first->type != GGML_TYPE_F32 || first->ne[2] != 1 || first->ne[3] != 1 || + !ggml_is_contiguous(first) || first->ne[0] <= 0 || first->ne[1] <= 0 || + first->ne[0] > (INT_MAX - CUDA_NEG_BLOCK_SIZE)/first->ne[1]) { return {}; } + const auto next_compute = [&](int index) { + while (index < graph->n_nodes && ggml_cuda_affine_unary_view_op(graph->nodes[index]->op)) { + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[index], true) || + (allocated && !ggml_cuda_affine_unary_range(graph->nodes[index], ctx.device, begin, end))) { return graph->n_nodes; } + ++index; + } + return index; + }; + ggml_cuda_ordered_mul_add_match match; + int indices[2*GGML_CUDA_ORDERED_MUL_ADD_MAX - 1]; + ggml_op ops[2*GGML_CUDA_ORDERED_MUL_ADD_MAX - 1]; + size_t step[2]{}; + int index = i; + int nodes = 0; + while (match.streams < GGML_CUDA_ORDERED_MUL_ADD_MAX && index < graph->n_nodes) { + ggml_tensor * product = graph->nodes[index]; + if (product->op != GGML_OP_MUL || product->type != GGML_TYPE_F32 || !ggml_are_same_layout(first, product) || + ggml_cuda_match_hc_affine_injection(ctx, graph, index, false).count || + ggml_cuda_match_affine_unary(ctx, graph, index, false).count) { break; } + bool valid = true; + for (int k = 0; k < 2; ++k) { + const ggml_tensor * tensor = product->src[k]; + const ggml_tensor * base = first->src[k]; + valid &= tensor && tensor->type == GGML_TYPE_F32 && tensor->view_src && + tensor->view_src->type == GGML_TYPE_F32 && ggml_cuda_affine_unary_prior_view(graph, i, tensor, true); + if (!valid) { break; } + valid &= ggml_are_same_layout(base, tensor) && tensor->view_src == base->view_src && + tensor->ne[2] == 1 && tensor->ne[3] == 1 && tensor->nb[0] == sizeof(float) && tensor->nb[1] % sizeof(float) == 0 && tensor->view_offs % sizeof(float) == 0 && + (k == 0 ? ggml_are_same_shape(first, tensor) : tensor->ne[0] == 1 && first->ne[1] % tensor->ne[1] == 0); + if (match.streams == 1) { + valid &= tensor->view_offs > base->view_offs; + step[k] = tensor->view_offs - base->view_offs; + } + if (match.streams) { + valid &= step[k] % sizeof(float) == 0 && step[k] <= (SIZE_MAX - base->view_offs)/match.streams && + tensor->view_offs == base->view_offs + size_t(match.streams)*step[k]; + } + } + if (!valid) { break; } + int add_i = match.streams ? next_compute(index + 1) : index; + ggml_tensor * sum = add_i < graph->n_nodes ? graph->nodes[add_i] : nullptr; + if (match.streams && (!sum || sum->op != GGML_OP_ADD || sum->type != GGML_TYPE_F32 || !ggml_are_same_layout(first, sum) || + !((sum->src[0] == product && sum->src[1] == match.sums[match.streams - 1]) || + (sum->src[1] == product && sum->src[0] == match.sums[match.streams - 1])))) { break; } + if (match.streams && add_i + 1 < graph->n_nodes && ggml_can_fuse(graph, add_i, {GGML_OP_ADD, GGML_OP_ADD})) { break; } + match.products[match.streams] = product; + match.sums[match.streams] = sum; + indices[nodes] = index; ops[nodes++] = GGML_OP_MUL; + if (match.streams) { indices[nodes] = add_i; ops[nodes++] = GGML_OP_ADD; } + ++match.streams; + match.count = add_i - i + 1; + index = next_compute(add_i + 1); + } + if (match.streams < 2 || !ggml_can_fuse_subgraph_ext(graph, indices, nodes, ops, indices, nodes)) { return {}; } + uintptr_t rb[2], re[2]; + for (int k = 0; k < 2; ++k) { + const ggml_tensor * root = first->src[k]->view_src; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(root, bytes) || (allocated && + (!ggml_cuda_affine_unary_range(root, ctx.device, rb[k], re[k]) || rb[k] % sizeof(float)))) { return {}; } + for (int j = 0; j < match.streams; ++j) { + uintptr_t begin, end; + if (allocated && !ggml_cuda_affine_unary_range(match.products[j]->src[k], ctx.device, begin, end)) { return {}; } + } + } + if (allocated) { + uintptr_t wb[2*GGML_CUDA_ORDERED_MUL_ADD_MAX - 1], we[2*GGML_CUDA_ORDERED_MUL_ADD_MAX - 1]; + for (int j = 0; j < nodes; ++j) { + if (!ggml_cuda_affine_unary_range(graph->nodes[indices[j]], ctx.device, wb[j], we[j]) || wb[j] % sizeof(float)) { return {}; } + for (int k = 0; k < 2; ++k) { if (wb[j] < re[k] && rb[k] < we[j]) { return {}; } } + for (int k = 0; k < j; ++k) { + if (wb[j] < we[k] && wb[k] < we[j] && (wb[j] != wb[k] || we[j] != we[k])) { return {}; } + } + // The next product must not overwrite the accumulator before ADD reads it. + if (j > 0 && ops[j] == GGML_OP_MUL && wb[j] == wb[j - 1]) { return {}; } + } + } + return match; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i, const ggml_tensor * shared_input, const char * quantized, int prepared_hc_post = -1, const ggml_cuda_hc_affine_injection_match * affine_match = nullptr, bool * affine_emit = nullptr, const void * prepared_src1 = nullptr) { @@ -6145,6 +6238,11 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_cuda_op_affine_unary(*cuda_ctx, affine.mul, affine.add, affine.unary, affine.post, affine.tail); return affine.count - 1; } + const auto ordered = ggml_cuda_match_ordered_mul_add(*cuda_ctx, cgraph, i); + if (ordered.count) { + ggml_cuda_op_ordered_mul_add(*cuda_ctx, ordered.products, ordered.sums, ordered.streams); + return ordered.count - 1; + } const auto binary = ggml_cuda_match_mul_add(*cuda_ctx, cgraph, i); if (binary.count) { ggml_cuda_op_mul_add(*cuda_ctx, binary.mul, binary.add); @@ -6873,6 +6971,18 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph continue; } + const auto ordered = ggml_cuda_match_ordered_mul_add(*cuda_ctx, cgraph, i, false); + if (ordered.count) { + ggml_tensor * terminal = ordered.sums[ordered.streams - 1]; + for (int j = 0; j < ordered.streams; ++j) { + for (ggml_tensor * read : {ordered.products[j]->src[0], ordered.products[j]->src[1]}) { + params->add_alloc_dep(params->user_data, read, terminal); + params->add_alloc_dep(params->user_data, read->view_src, terminal); + } + } + i += ordered.count - 1; + continue; + } const auto binary = ggml_cuda_match_mul_add(*cuda_ctx, cgraph, i, false); if (binary.count) { for (ggml_tensor * read : {binary.mul->src[0], binary.mul->src[1], binary.add->src[binary.add->src[0] == binary.mul ? 1 : 0]}) { diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index fc02803fd22..7d3907d0633 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -494,6 +494,63 @@ void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * ggml_cuda_kernel_launch(repeat_mul_add_kernel, launch, binary, mul->src[0] == repeat ? 0 : 1, add->src[0] == mul); } +struct ordered_mul_add_args { + const char * input[2]; + size_t row_stride[2]; + size_t stream_stride[2]; + float * products[GGML_CUDA_ORDERED_MUL_ADD_MAX]; + float * sums[GGML_CUDA_ORDERED_MUL_ADD_MAX]; + bool product_first[GGML_CUDA_ORDERED_MUL_ADD_MAX]; + uint3 columns; + uint3 coefficient_rows; + int elements; + int count; +}; + +static __global__ void ordered_mul_add_kernel(ordered_mul_add_args args) { + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= args.elements) { return; } + const uint2 pos = fast_div_modulo(i, args.columns); + const size_t x_row = size_t(pos.x)*args.row_stride[0] + size_t(pos.y)*sizeof(float); + const size_t w_row = size_t(fastmodulo(pos.x, args.coefficient_rows))*args.row_stride[1]; + ggml_cuda_pdl_sync(); + float sum = 0.0f; + for (int j = 0; j < args.count; ++j) { + const float x = *(const float *) (args.input[0] + x_row + size_t(j)*args.stream_stride[0]); + const float w = *(const float *) (args.input[1] + w_row + size_t(j)*args.stream_stride[1]); + const float product = __fmul_rn(x, w); + args.products[j][i] = product; + if (j == 0) { sum = product; } + else { + sum = args.product_first[j] ? __fadd_rn(product, sum) : __fadd_rn(sum, product); + args.sums[j][i] = sum; + } + } +} + +void ggml_cuda_op_ordered_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * const * products, ggml_tensor * const * sums, int count) { + ordered_mul_add_args args{}; + for (int k = 0; k < 2; ++k) { + const ggml_tensor * first = products[0]->src[k]; + args.input[k] = (const char *) first->data; + args.row_stride[k] = first->nb[1]; + args.stream_stride[k] = products[1]->src[k]->view_offs - first->view_offs; + } + for (int j = 0; j < count; ++j) { + args.products[j] = (float *) products[j]->data; + args.sums[j] = (float *) sums[j]->data; + args.product_first[j] = j && sums[j]->src[0] == products[j]; + } + args.columns = init_fastdiv_values(products[0]->ne[0]); + args.coefficient_rows = init_fastdiv_values(products[0]->src[1]->ne[1]); + args.elements = ggml_nelements(products[0]); + args.count = count; + const int blocks = (args.elements + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(ordered_mul_add_kernel, launch, args); +} + static void affine_unary_tail_cuda(ggml_backend_cuda_context & ctx, affine_unary_args args, ggml_tensor * tail, const void * src = nullptr, ggml_type src_type = GGML_TYPE_F32, ggml_tensor * mm = nullptr, int offset = 0) { GGML_ASSERT(args.output[3] && tail && (tail->type == GGML_TYPE_F32 || tail->type == GGML_TYPE_BF16)); diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 909ac2a47b3..635d774ab51 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -96,6 +96,9 @@ void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add); void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * repeat, ggml_tensor * mul, ggml_tensor * add); +constexpr int GGML_CUDA_ORDERED_MUL_ADD_MAX = 128; +void ggml_cuda_op_ordered_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * const * products, ggml_tensor * const * sums, int count); + void ggml_cuda_op_affine_unary(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * unary, ggml_tensor * post, ggml_tensor * tail = nullptr); void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 09fd82c804a..460cbdcd5b8 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4428,24 +4428,28 @@ struct test_dsv4_hc_pre : public test_dsv4_hc { const bool gated; const ggml_type projection_type; const bool computed; + const bool ordered; + const bool swap_add; + std::vector stages; ggml_tensor * projection = nullptr; ggml_tensor * projection_weights = nullptr; ggml_tensor * projection_input = nullptr; std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); - return "DSV4_HC_PRE"; + return ordered ? "ORDERED_MUL_ADD" : "DSV4_HC_PRE"; } std::string vars() override { - return VARS_TO_STR4(n_embd, n_hc, n_tokens, gated) + (projection_type == GGML_TYPE_COUNT ? "" : "," + VARS_TO_STR2(projection_type, computed)); + return VARS_TO_STR4(n_embd, n_hc, n_tokens, gated) + (projection_type == GGML_TYPE_COUNT ? "" : "," + VARS_TO_STR2(projection_type, computed)) + (ordered ? "," + VARS_TO_STR2(ordered, swap_add) : ""); } - bool run_whole_graph() override { return projection_type != GGML_TYPE_COUNT; } - std::vector fusion_test_nodes() override { return projection ? std::vector{projection, out} : std::vector{}; } + bool run_whole_graph() override { return ordered || projection_type != GGML_TYPE_COUNT; } + std::vector fusion_test_nodes() override { return ordered ? stages : projection ? std::vector{projection, out} : std::vector{}; } + double max_nmse_err() override { return ordered ? 1e-7 : test_dsv4_hc::max_nmse_err(); } - test_dsv4_hc_pre(int64_t n_embd = 31, int64_t n_hc = 4, int64_t n_tokens = 17, bool gated = false, ggml_type projection_type = GGML_TYPE_COUNT, bool computed = false) - : n_embd(n_embd), n_hc(n_hc), n_tokens(n_tokens), gated(gated), projection_type(projection_type), computed(computed) {} + test_dsv4_hc_pre(int64_t n_embd = 31, int64_t n_hc = 4, int64_t n_tokens = 17, bool gated = false, ggml_type projection_type = GGML_TYPE_COUNT, bool computed = false, bool ordered = false, bool swap_add = false) + : n_embd(n_embd), n_hc(n_hc), n_tokens(n_tokens), gated(gated), projection_type(projection_type), computed(computed), ordered(ordered), swap_add(swap_add) {} void initialize_tensors(ggml_context * ctx) override { test_dsv4_hc::initialize_tensors(ctx); @@ -4492,7 +4496,24 @@ struct test_dsv4_hc_pre : public test_dsv4_hc { ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_hc, n_tokens); ggml_set_name(weights, "weights"); - out = ggml_dsv4_hc_pre(ctx, x, weights); + if (ordered) { + stages.clear(); + out = nullptr; + for (int64_t stream = 0; stream < n_hc; ++stream) { + ggml_tensor * xv = ggml_view_2d(ctx, x, n_embd, n_tokens, x->nb[2], stream*x->nb[1]); + ggml_tensor * wv = ggml_view_2d(ctx, weights, 1, n_tokens, weights->nb[1], stream*weights->nb[0]); + ggml_tensor * product = ggml_mul(ctx, xv, wv); + ggml_set_output(product); + stages.push_back(product); + if (out) { + out = ggml_add(ctx, swap_add ? product : out, swap_add ? out : product); + ggml_set_output(out); + stages.push_back(out); + } else { out = product; } + } + } else { + out = ggml_dsv4_hc_pre(ctx, x, weights); + } } ggml_set_name(out, "out"); return out; @@ -9408,6 +9429,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_comb(n_tokens, 20)); } + for (bool swap_add : {false, true}) { + for (int64_t streams : {1, 4, 7, 16, 24, 30}) { + test_cases.emplace_back(new test_dsv4_hc_pre(31, streams, 3, false, GGML_TYPE_COUNT, false, true, swap_add)); + } + test_cases.emplace_back(new test_dsv4_hc_pre(512, 4, 32, false, GGML_TYPE_COUNT, false, true, swap_add)); + } test_cases.emplace_back(new test_dsv4_hc_pre(1, 4, 1)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 4, 17)); test_cases.emplace_back(new test_dsv4_hc_pre(128, 4, 257)); From 5fa22f10e869d65bb35c710103a49c32bafa6079 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 19:06:54 -0400 Subject: [PATCH 15/27] cuda : fuse deferred REPEAT and MUL ADD Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 86 +++++++++++++++++++++++++++++++-- tests/test-backend-ops.cpp | 27 +++++++++-- 2 files changed, 107 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6757075931e..c63ffbb56e1 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5248,13 +5248,14 @@ struct ggml_cuda_repeat_mul_add_match { ggml_tensor * repeat = nullptr; ggml_cuda_affine_unary_match binary; int count = 0; + int mul_i = 0; }; -static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { +static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true, int consumer = -1) { ggml_tensor * repeat = graph->nodes[i]; if (repeat->op != GGML_OP_REPEAT || !repeat->src[0] || repeat->type != repeat->src[0]->type || (repeat->type != GGML_TYPE_F32 && repeat->type != GGML_TYPE_F16)) { return {}; } - int m = i + 1; + int m = consumer < 0 ? i + 1 : consumer; while (m < graph->n_nodes && ggml_cuda_affine_unary_view_op(graph->nodes[m]->op)) { uintptr_t begin, end; if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[m], true) || @@ -5292,7 +5293,58 @@ static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backen } } } - return {repeat, binary, indices[2] - i + 1}; + return {repeat, binary, indices[2] - i + 1, m}; +} + + +static ggml_cuda_repeat_mul_add_match ggml_cuda_match_deferred_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * repeat = graph->nodes[i]; + if (repeat->op != GGML_OP_REPEAT || !repeat->src[0]) { return {}; } + int m = i + 1; + bool compute = false; + for (; m < graph->n_nodes; ++m) { + const ggml_tensor * node = graph->nodes[m]; + bool used = node->view_src == repeat; + for (const ggml_tensor * src : node->src) { used |= src && (src == repeat || src->view_src == repeat); } + if (used) { break; } + compute |= !ggml_cuda_affine_unary_view_op(node->op); + } + if (!compute || m == graph->n_nodes) { return {}; } + const auto match = ggml_cuda_match_repeat_mul_add(ctx, graph, i, allocated, m); + if (!match.count) { return {}; } + // Keep earlier producer and unary fusions across the MUL boundary. + const ggml_tensor * prior = graph->nodes[m - 1]; + if (!ggml_cuda_affine_unary_view_op(prior->op) && prior->op != GGML_OP_SCALE && prior->op != GGML_OP_CONT) { return {}; } + ggml_cuda_moe_weighted_reduction_match reduction; + if (ggml_cuda_match_moe_weighted_reduction(graph, m, reduction)) { return {}; } + for (int j = i + 1; j < m; ++j) { + const ggml_tensor * node = graph->nodes[j]; + switch (node->op) { + case GGML_OP_VIEW: case GGML_OP_RESHAPE: case GGML_OP_PERMUTE: case GGML_OP_TRANSPOSE: + case GGML_OP_MUL_MAT: case GGML_OP_SCALE: case GGML_OP_UNARY: case GGML_OP_MUL: + case GGML_OP_SUM_ROWS: case GGML_OP_SUB: case GGML_OP_CONT: break; + default: return {}; + } + } + if (allocated) { + const ggml_tensor * input = repeat->src[0]->view_src ? repeat->src[0]->view_src : repeat->src[0]; + uintptr_t ib, ie, wb, we; + if (!ggml_cuda_affine_unary_range(input, ctx.device, ib, ie) || !ggml_cuda_affine_unary_range(repeat, ctx.device, wb, we)) { return {}; } + for (const ggml_tensor * write : {repeat, match.binary.mul, match.binary.add}) { + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_range(write, ctx.device, begin, end) || (begin < ie && ib < end)) { return {}; } + } + for (int j = i + 1; j < m; ++j) { + const ggml_tensor * node = graph->nodes[j]; + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_range(node, ctx.device, begin, end)) { return {}; } + if (!ggml_cuda_affine_unary_view_op(node->op) && ((begin < ie && ib < end) || (begin < we && wb < end))) { return {}; } + for (const ggml_tensor * src : node->src) { + if (src && (!ggml_cuda_affine_unary_range(src, ctx.device, begin, end) || (begin < we && wb < end))) { return {}; } + } + } + } + return match; } struct ggml_cuda_ordered_mul_add_match { @@ -6474,6 +6526,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } }; + std::unordered_map repeat_emits; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -6516,6 +6569,23 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud continue; } + if (!disable_reuse && cuda_ctx->curr_stream_no == 0 && stream_ctx.concurrent_events.empty()) { + if (node->op == GGML_OP_REPEAT) { + const auto deferred = ggml_cuda_match_deferred_repeat_mul_add(*cuda_ctx, cgraph, i); + if (deferred.count) { + repeat_emits.emplace(deferred.binary.mul, deferred); + continue; + } + } + const auto emit = repeat_emits.find(node); + if (emit != repeat_emits.end()) { + const auto & match = emit->second; + ggml_cuda_op_repeat_mul_add(*cuda_ctx, match.repeat, match.binary.mul, match.binary.add); + i += match.binary.count - 1; + continue; + } + } + const int group = reuse.nodes.empty() ? -1 : reuse.nodes[i]; prepare_shared(i, false); if (group >= 0) { prepare_group(group); } @@ -6936,6 +7006,16 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto deferred = ggml_cuda_match_deferred_repeat_mul_add(*cuda_ctx, cgraph, i, false); + if (deferred.count) { + ggml_tensor * read = deferred.repeat->src[0]; + params->add_alloc_dep(params->user_data, read, deferred.binary.add); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, deferred.binary.add); } + params->add_alloc_dep(params->user_data, deferred.repeat, deferred.binary.add); + params->add_alloc_dep(params->user_data, deferred.binary.mul->src[0], deferred.binary.add); + params->add_alloc_dep(params->user_data, deferred.binary.mul->src[1], deferred.binary.add); + } + const auto repeated = ggml_cuda_match_repeat_mul_add(*cuda_ctx, cgraph, i, false); if (repeated.count) { for (ggml_tensor * read : {repeated.repeat->src[0], repeated.repeat, repeated.binary.mul->src[0], repeated.binary.mul->src[1], repeated.binary.add->src[repeated.binary.add->src[0] == repeated.binary.mul ? 1 : 0]}) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 460cbdcd5b8..9043a708d8c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3299,13 +3299,14 @@ struct test_bin_bcast : public test_case { const bool add_after; const bool swap_add; const int repeat_input; + const bool coefficient_gap; std::vector stages; bool run_whole_graph() override { return nf > 1 || add_after; } std::vector fusion_test_nodes() override { return stages; } std::string op_desc(ggml_tensor * t) override { - return repeat_input >= 0 ? "REPEAT_MUL_ADD" : add_after ? "MUL_ADD" : test_case::op_desc(t); + return coefficient_gap ? "DEFERRED_REPEAT_MUL_ADD" : repeat_input >= 0 ? "REPEAT_MUL_ADD" : add_after ? "MUL_ADD" : test_case::op_desc(t); } std::string vars() override { @@ -3315,6 +3316,7 @@ struct test_bin_bcast : public test_case { } if (add_after) { s += "," + VARS_TO_STR2(add_after, swap_add); } if (repeat_input >= 0) { s += "," + VAR_TO_STR(repeat_input); } + if (coefficient_gap) { s += "," + VAR_TO_STR(coefficient_gap); } return s; } @@ -3327,8 +3329,8 @@ struct test_bin_bcast : public test_case { std::array nr = {1, 2, 1, 1}, int nf = 1, bool perm1 = false, bool src_overlap = false, - ggml_type type_b = GGML_TYPE_COUNT, bool add_after = false, bool swap_add = false, int repeat_input = -1) - : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b), add_after(add_after), swap_add(swap_add), repeat_input(repeat_input) {} + ggml_type type_b = GGML_TYPE_COUNT, bool add_after = false, bool swap_add = false, int repeat_input = -1, bool coefficient_gap = false) + : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1), src_overlap(src_overlap), type_b(type_b), add_after(add_after), swap_add(swap_add), repeat_input(repeat_input), coefficient_gap(coefficient_gap) {} ggml_tensor * build_graph(ggml_context * ctx) override { GGML_ASSERT(nf <= 16); @@ -3364,6 +3366,16 @@ struct test_bin_bcast : public test_case { else b[0] = repeated; } + if (coefficient_gap) { + GGML_ASSERT(repeat_input == 0 && tb == GGML_TYPE_F32 && !perm1 && !src_overlap); + b[0] = ggml_scale(ctx, b[0], 0.25f); + ggml_set_output(b[0]); stages.push_back(b[0]); + b[0] = ggml_sigmoid(ctx, b[0]); + ggml_set_output(b[0]); stages.push_back(b[0]); + b[0] = ggml_scale(ctx, b[0], 2.0f); + ggml_set_output(b[0]); stages.push_back(b[0]); + } + // The backward pass supports broadcasting only for GGML_ADD: const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; if (grad_supported) { @@ -3417,6 +3429,9 @@ struct test_bin_bcast : public test_case { } double max_nmse_err() override { + if (coefficient_gap) { + return 1e-7; + } if (op == ggml_add && nf > 1) { // Fused ADDs can keep FP32 intermediates while the CPU rounds each ADD to FP16/BF16. if (type == GGML_TYPE_F16) { @@ -10091,6 +10106,12 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { + for (bool swap : {false, true}) { + test_cases.emplace_back(new test_bin_bcast(ggml_mul, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, false, false, GGML_TYPE_F32, true, swap, 0, true)); + } + } + auto add_test_bin_bcast = [&](ggml_type type, std::array ne, std::array nr, bool perm1 = false, bool src_overlap = false) { for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) { test_cases.emplace_back(new test_bin_bcast(op, type, ne, nr, 1, perm1, src_overlap)); From 76479e8521cafecbe0461ca1899b7d928e515187 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 20:35:00 -0400 Subject: [PATCH 16/27] cuda : fuse products into RMS normalization Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 116 +++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/norm.cu | 117 +++++++++++++++++++++++++++++++- ggml/src/ggml-cuda/norm.cuh | 3 + tests/test-backend-ops.cpp | 35 ++++++++-- 4 files changed, 264 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c63ffbb56e1..bf85951977d 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5244,6 +5244,97 @@ static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_ return {mm, affine, count}; } +struct ggml_cuda_pre_mul_norm_match { + ggml_tensor * repeat = nullptr; + ggml_tensor * product = nullptr; + ggml_tensor * norm = nullptr; + ggml_tensor * weighted = nullptr; + int norm_i = -1; + int count = 0; +}; + +static ggml_cuda_pre_mul_norm_match ggml_cuda_match_pre_mul_norm(ggml_backend_cuda_context & ctx, + ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * repeat = graph->nodes[i]->op == GGML_OP_REPEAT ? graph->nodes[i] : nullptr; + const int m = i + int(repeat != nullptr); + if (m >= graph->n_nodes) { return {}; } + ggml_tensor * product = graph->nodes[m]; + if (product->op != GGML_OP_MUL || !product->src[0] || !product->src[1] || product->type != GGML_TYPE_F32 || + !ggml_is_contiguous(product) || ggml_nelements(product) <= 0 || ggml_nelements(product) > INT_MAX) { return {}; } + if (repeat && (!repeat->src[0] || !ggml_are_same_shape(repeat, product) || !ggml_is_contiguous(repeat) || + (product->src[0] == repeat) == (product->src[1] == repeat))) { return {}; } + int n = m + 1; + while (n < graph->n_nodes && (graph->nodes[n]->op == GGML_OP_VIEW || graph->nodes[n]->op == GGML_OP_RESHAPE)) { + const ggml_tensor * view = graph->nodes[n]; + if (view->view_src == product) { + if (view->view_offs || !ggml_is_contiguous(view) || ggml_nelements(view) != ggml_nelements(product)) { return {}; } + } else if (!ggml_cuda_affine_unary_prior_view(graph, i, view)) { return {}; } + ++n; + } + if (n >= graph->n_nodes) { return {}; } + ggml_tensor * norm = graph->nodes[n]; + if (norm->op != GGML_OP_RMS_NORM || !norm->src[0] || norm->type != GGML_TYPE_F32 || norm->src[0]->type != GGML_TYPE_F32 || + (norm->src[0] != product && (norm->src[0]->view_src != product || norm->src[0]->view_offs)) || + !ggml_is_contiguous(norm->src[0]) || !ggml_is_contiguous(norm) || ggml_nelements(norm) != ggml_nelements(product) || + !ggml_are_same_shape(norm, norm->src[0]) || norm->ne[0] > INT_MAX - 1023 || norm->ne[2] > 65535 || norm->ne[3] > 65535) { return {}; } + ggml_tensor * weighted = nullptr; + int end = n + 1; + while (end < graph->n_nodes && ggml_cuda_affine_unary_view_op(graph->nodes[end]->op)) { + if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[end], true)) { return {}; } + ++end; + } + if (end < graph->n_nodes && graph->nodes[end]->op == GGML_OP_MUL && + (graph->nodes[end]->src[0] == norm) != (graph->nodes[end]->src[1] == norm)) { + weighted = graph->nodes[end++]; + if (!ggml_are_same_shape(norm, weighted) || !ggml_is_contiguous(weighted)) { return {}; } + } + // Leave longer normalization fusions to the existing dispatch. + if (end < graph->n_nodes && (graph->nodes[end]->op == GGML_OP_ROPE || graph->nodes[end]->op == GGML_OP_ADD || graph->nodes[end]->op == GGML_OP_SCALE)) { return {}; } + if (!(ggml_get_op_params_f32(norm, 0) >= 0.0f)) { return {}; } + const ggml_tensor * writes[] = {product, norm, repeat, weighted}; + for (const ggml_tensor * write : writes) { + if (write && (write->type != GGML_TYPE_F32 || !(write->flags & GGML_TENSOR_FLAG_COMPUTE))) { return {}; } + } + std::vector reads; + for (int j = 0; j < 2; ++j) { reads.push_back(product->src[j] == repeat ? repeat->src[0] : product->src[j]); } + if (weighted) { reads.push_back(weighted->src[weighted->src[0] == norm ? 1 : 0]); } + for (size_t j = 0; j < reads.size(); ++j) { + const ggml_tensor * read = reads[j]; + const ggml_tensor * shape = j < 2 ? product : norm; + size_t bytes; + if (!read || read->type != GGML_TYPE_F32 || !ggml_cuda_affine_unary_bytes(read, bytes)) { return {}; } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (read->ne[d] > INT_MAX || read->nb[d] % sizeof(float) || read->nb[d]/sizeof(float) > UINT32_MAX || shape->ne[d] % read->ne[d]) { return {}; } + } + } + std::vector indices; + std::vector ops; + for (int j = i; j < end; ++j) { indices.push_back(j); ops.push_back(graph->nodes[j]->op); } + if (!ggml_can_fuse_subgraph_ext(graph, indices.data(), indices.size(), ops.data(), indices.data(), indices.size())) { return {}; } + if (allocated) { + for (int j = i; j < end; ++j) { + uintptr_t begin, stop; + if (!ggml_cuda_affine_unary_range(graph->nodes[j], ctx.device, begin, stop) || begin % sizeof(float)) { return {}; } + } + std::vector writes = {product, norm}; + if (repeat) { writes.push_back(repeat); } + if (weighted) { writes.push_back(weighted); } + std::vector wb(writes.size()), we(writes.size()); + for (size_t j = 0; j < writes.size(); ++j) { + if (writes[j]->type != GGML_TYPE_F32 || !(writes[j]->flags & GGML_TENSOR_FLAG_COMPUTE) || + !ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j])) { return {}; } + for (size_t k = 0; k < j; ++k) { if (wb[j] < we[k] && wb[k] < we[j]) { return {}; } } + for (const ggml_tensor * read : reads) { + const ggml_tensor * root = read->view_src ? read->view_src : read; + uintptr_t rb, re; + if (!ggml_cuda_affine_unary_range(read, ctx.device, rb, re) || rb % sizeof(float) || + !ggml_cuda_affine_unary_range(root, ctx.device, rb, re) || (wb[j] < re && rb < we[j])) { return {}; } + } + } + } + return {repeat, product, norm, weighted, n, end - i}; +} + struct ggml_cuda_repeat_mul_add_match { ggml_tensor * repeat = nullptr; ggml_cuda_affine_unary_match binary; @@ -6777,6 +6868,16 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } } + if (!disable_reuse && cuda_ctx->curr_stream_no == 0 && stream_ctx.concurrent_events.empty()) { + const auto pre = ggml_cuda_match_pre_mul_norm(*cuda_ctx, cgraph, i); + if (pre.count && (norm_emits.empty() || !norm_emits[pre.norm_i].norm) && + (q8_emits.empty() || !q8_emits[pre.norm_i].norm) && (mmq_emits.empty() || !mmq_emits[pre.norm_i].norm)) { + ggml_cuda_op_rms_norm_pre_mul(*cuda_ctx, pre.product, pre.norm, pre.weighted, pre.repeat); + i += pre.count - 1; + continue; + } + } + int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i, shared_input, quantized, prepared_hc_post, affine_match, nullptr, prepared_src1); if (nodes_to_skip != 0) { @@ -7006,6 +7107,21 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto pre = ggml_cuda_match_pre_mul_norm(*cuda_ctx, cgraph, i, false); + if (pre.count) { + const int last = i + pre.count - 1; + add_alloc_deps(i, last); + for (int j = i; j <= last; ++j) { + for (ggml_tensor * read : cgraph->nodes[j]->src) { + if (!read) { continue; } + params->add_alloc_dep(params->user_data, read, cgraph->nodes[last]); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, cgraph->nodes[last]); } + } + } + i = last; + continue; + } + const auto deferred = ggml_cuda_match_deferred_repeat_mul_add(*cuda_ctx, cgraph, i, false); if (deferred.count) { ggml_tensor * read = deferred.repeat->src[0]; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c5f00f40ebd..f5af15de61c 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -313,7 +313,11 @@ struct ggml_cuda_norm_mxfp4_store { } }; -template +struct ggml_cuda_norm_load { + __device__ __forceinline__ float operator()(const float * x, int col) const { return x[col]; } +}; + +template static __device__ __forceinline__ void rms_norm_f32_impl(const float * x, float * dst, const int ncols, @@ -338,7 +342,8 @@ static __device__ __forceinline__ void rms_norm_f32_impl(const float * x, const uint3 add_nrows_packed = make_uint3(0, 0, 0), const uint3 add_nchannels_packed = make_uint3(0, 0, 0), const uint3 add_nsamples_packed = make_uint3(0, 0, 0), - const float scale_out = 1.0f) { + const float scale_out = 1.0f, + const Read read = {}) { ggml_cuda_pdl_lc(); const int nrows = gridDim.x; const int nchannels = gridDim.y; @@ -373,7 +378,7 @@ static __device__ __forceinline__ void rms_norm_f32_impl(const float * x, ggml_cuda_pdl_sync(); for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; + const float xi = read(x, col); tmp += xi * xi; } @@ -421,6 +426,77 @@ static __device__ __forceinline__ void rms_norm_f32_impl(const float * x, } } +struct ggml_cuda_norm_mul_input { + const float * data; + uint3 ne[GGML_MAX_DIMS]; + uint32_t stride[GGML_MAX_DIMS]; + + __device__ __forceinline__ float load(const uint32_t * coords) const { + uint64_t offset = 0; +#pragma unroll + for (int d = 0; d < GGML_MAX_DIMS; ++d) { offset += uint64_t(fastmodulo(coords[d], ne[d])) * stride[d]; } + return data[offset]; + } +}; + +struct ggml_cuda_norm_pre_mul_load { + ggml_cuda_norm_mul_input inputs[2]; + uint3 ne[3]; + uint32_t offset; + float * product; + float * repeat; + int repeated; + + __device__ __forceinline__ float operator()(const float * x, int col) const { + const uint32_t index = offset + col; + uint32_t coords[GGML_MAX_DIMS]; + uint32_t remainder = index; +#pragma unroll + for (int d = 0; d < 3; ++d) { + const uint2 pair = fast_div_modulo(remainder, ne[d]); + coords[d] = pair.y; remainder = pair.x; + } + coords[3] = remainder; + const float a = inputs[0].load(coords); + const float b = inputs[1].load(coords); + if (repeat) { repeat[index] = repeated == 0 ? a : b; } + const float value = __fmul_rn(a, b); + product[index] = value; + GGML_UNUSED(x); + return value; + } +}; + +struct ggml_cuda_norm_pre_mul_store { + static constexpr int width = 1; + ggml_cuda_norm_mul_input weight; + float * weighted; + uint32_t offset; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float value) const { + dst[col] = value; + if (weighted) { weighted[offset + col] = __fmul_rn(value, weight.data[uint64_t(fastmodulo(col, weight.ne[0])) * weight.stride[0]]); } + GGML_UNUSED(base); + } +}; + +template +static __global__ __launch_bounds__(block_size) void rms_norm_pre_mul_f32( + ggml_cuda_norm_pre_mul_load read, ggml_cuda_norm_pre_mul_store write, + float * dst, int ncols, float eps) { + const uint32_t row = (blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x; + read.offset = write.offset = row * ncols; + if (write.weighted) { + write.weight.data += uint64_t(fastmodulo(blockIdx.x, write.weight.ne[1])) * write.weight.stride[1] + + uint64_t(fastmodulo(blockIdx.y, write.weight.ne[2])) * write.weight.stride[2] + + uint64_t(fastmodulo(blockIdx.z, write.weight.ne[3])) * write.weight.stride[3]; + } + rms_norm_f32_impl(read.product, dst, ncols, ncols, + int64_t(ncols)*gridDim.x, int64_t(ncols)*gridDim.x*gridDim.y, eps, write, + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f, read); +} + template static __global__ void rms_norm_f32(const float * x, float * dst, @@ -960,6 +1036,41 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } +void ggml_cuda_op_rms_norm_pre_mul(ggml_backend_cuda_context & ctx, ggml_tensor * product, + ggml_tensor * norm, ggml_tensor * weighted, ggml_tensor * repeat) { + const auto input = [](const ggml_tensor * tensor) { + ggml_cuda_norm_mul_input value{}; + if (!tensor) { return value; } + value.data = (const float *) tensor->data; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + value.ne[d] = init_fastdiv_values(tensor->ne[d]); + value.stride[d] = tensor->nb[d]/sizeof(float); + } + return value; + }; + ggml_cuda_norm_pre_mul_load read{}; + for (int j = 0; j < 2; ++j) { + const ggml_tensor * src = product->src[j]; + if (src == repeat) { src = repeat->src[0]; read.repeated = j; } + read.inputs[j] = input(src); + } + for (int d = 0; d < 3; ++d) { read.ne[d] = init_fastdiv_values(product->ne[d]); } + read.product = (float *) product->data; + read.repeat = repeat ? (float *) repeat->data : nullptr; + ggml_cuda_norm_pre_mul_store write{}; + if (weighted) { + write.weight = input(weighted->src[weighted->src[0] == norm ? 1 : 0]); + write.weighted = (float *) weighted->data; + } + const int ncols = norm->ne[0]; + const dim3 blocks(norm->ne[1], norm->ne[2], norm->ne[3]); + const float eps = ggml_get_op_params_f32(norm, 0); + GGML_ASSERT(eps >= 0.0f); + const ggml_cuda_kernel_launch_params params = {blocks, dim3(ncols < 1024 ? 256 : 1024), 32*sizeof(float), ctx.stream()}; + if (ncols < 1024) { ggml_cuda_kernel_launch(rms_norm_pre_mul_f32<256>, params, read, write, (float *) norm->data, ncols, eps); } + else { ggml_cuda_kernel_launch(rms_norm_pre_mul_f32<1024>, params, read, write, (float *) norm->data, ncols, eps); } +} + void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 2c80cc8d429..21d3ebb4487 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -6,6 +6,9 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_rms_norm_pre_mul(ggml_backend_cuda_context & ctx, ggml_tensor * product, + ggml_tensor * norm, ggml_tensor * weighted, ggml_tensor * repeat); + void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 9043a708d8c..25cfeae9a29 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3837,26 +3837,47 @@ struct test_rms_norm_mul_add : public test_case { const bool post_mul; const bool alias_rms_input; const bool weight_broadcast; + const bool pre_mul; + const bool repeat; + std::vector stages; + + std::vector fusion_test_nodes() override { return stages; } std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); - return "RMS_NORM_MUL_ADD"; + return pre_mul ? "MUL_RMS_NORM" : "RMS_NORM_MUL_ADD"; } bool run_whole_graph() override { return true; } std::string vars() override { - return VARS_TO_STR8(type, ne, eps, broadcast, multi_add, post_mul, alias_rms_input, weight_broadcast); + return VARS_TO_STR8(type, ne, eps, broadcast, multi_add, post_mul, alias_rms_input, weight_broadcast) + (pre_mul ? "," + VARS_TO_STR2(pre_mul, repeat) : ""); } test_rms_norm_mul_add(ggml_type type = GGML_TYPE_F32, std::array ne = {64, 5, 4, 3}, float eps = 1e-6f, bool broadcast = false, bool multi_add = false, bool post_mul = false, - bool alias_rms_input = false, bool weight_broadcast = false) + bool alias_rms_input = false, bool weight_broadcast = false, bool pre_mul = false, bool repeat = false) : type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add), post_mul(post_mul), - alias_rms_input(alias_rms_input), weight_broadcast(weight_broadcast) {} + alias_rms_input(alias_rms_input), weight_broadcast(weight_broadcast), pre_mul(pre_mul), repeat(repeat) {} ggml_tensor * build_graph(ggml_context * ctx) override { + stages.clear(); + if (pre_mul) { + GGML_ASSERT(type == GGML_TYPE_F32 && !broadcast && !multi_add && !post_mul && !alias_rms_input); + ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], repeat ? 1 : ne[1], ne[2], ne[3]); + ggml_tensor * b = ggml_new_tensor_4d(ctx, type, 1, ne[1], ne[2], ne[3]); + ggml_tensor * w = ggml_new_tensor_2d(ctx, type, ne[0], ne[1]); + for (ggml_tensor * input : {a, b, w}) { ggml_set_param(input); } + if (repeat) { a = ggml_repeat_4d(ctx, a, ne[0], ne[1], ne[2], ne[3]); stages.push_back(a); } + ggml_tensor * product = ggml_mul(ctx, a, b); + ggml_tensor * grouped = ggml_reshape_3d(ctx, product, ne[0], ne[1], ne[2]*ne[3]); + ggml_tensor * norm = ggml_rms_norm(ctx, grouped, eps); + ggml_tensor * out = ggml_mul(ctx, norm, w); + for (ggml_tensor * stage : {product, norm, out}) { stages.push_back(stage); } + for (ggml_tensor * stage : stages) { ggml_set_output(stage); } + return out; + } std::array broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4}; ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data()); @@ -10260,6 +10281,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_add_add(GGML_TYPE_F16, GGML_TYPE_F32, { n, 5, 4, 3 }, true, false)); } + for (int64_t n : {257, 1536}) { + for (bool repeat : {false, true}) { + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {n, 4, 3, 2}, 1e-6f, false, false, false, false, true, true, repeat)); + } + } + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 1536, 1, 1, 1 }, 1e-6f, false, false, true)); test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 1, 1 }, 1e-6f, false, false, true)); test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 3, 2 }, 1e-6f, false, false, true)); From 4af8f37904be05d6412e5c2c1963018b933e0447 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 21:20:45 -0400 Subject: [PATCH 17/27] cuda : defer repeat into product RMS normalization Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 167 ++++++++++++++++++++++++-------- 1 file changed, 124 insertions(+), 43 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index bf85951977d..4f6d3fa9797 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5244,6 +5244,68 @@ static ggml_cuda_cublas_affine_match ggml_cuda_match_cublas_affine(ggml_backend_ return {mm, affine, count}; } +static int ggml_cuda_repeat_first_consumer(ggml_cgraph * graph, int i, bool & compute) { + const ggml_tensor * repeat = graph->nodes[i]; + int m = i + 1; + compute = false; + for (; m < graph->n_nodes; ++m) { + const ggml_tensor * node = graph->nodes[m]; + bool used = node->view_src == repeat; + for (const ggml_tensor * src : node->src) { used |= src && (src == repeat || src->view_src == repeat); } + if (used) { break; } + compute |= !ggml_cuda_affine_unary_view_op(node->op); + } + return m; +} + +static bool ggml_cuda_repeat_interval_ok(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, int m, + std::initializer_list writes, bool allocated, bool normalization = false) { + const ggml_tensor * repeat = graph->nodes[i]; + for (int j = i + 1; j < m; ++j) { + const ggml_tensor * node = graph->nodes[j]; + switch (node->op) { + case GGML_OP_VIEW: case GGML_OP_RESHAPE: case GGML_OP_PERMUTE: case GGML_OP_TRANSPOSE: + case GGML_OP_MUL_MAT: case GGML_OP_SCALE: case GGML_OP_UNARY: case GGML_OP_MUL: + case GGML_OP_SUM_ROWS: case GGML_OP_SUB: case GGML_OP_CONT: break; + case GGML_OP_NONE: + if (!normalization || (node->flags & GGML_TENSOR_FLAG_COMPUTE)) { return false; } + break; + case GGML_OP_RMS_NORM: case GGML_OP_SQRT: case GGML_OP_CLAMP: case GGML_OP_ADD: + if (!normalization) { return false; } + break; + default: return false; + } + } + if (allocated) { + const ggml_tensor * input = repeat->src[0]->view_src ? repeat->src[0]->view_src : repeat->src[0]; + uintptr_t ib, ie, wb, we; + if (!ggml_cuda_affine_unary_range(input, ctx.device, ib, ie) || !ggml_cuda_affine_unary_range(repeat, ctx.device, wb, we)) { return false; } + for (const ggml_tensor * write : writes) { + if (!write) { continue; } + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_range(write, ctx.device, begin, end) || (begin < ie && ib < end)) { return false; } + } + for (int j = i + 1; j < m; ++j) { + const ggml_tensor * node = graph->nodes[j]; + uintptr_t begin, end; + if (!ggml_cuda_affine_unary_range(node, ctx.device, begin, end)) { return false; } + if (normalization && node->op == GGML_OP_NONE) { + const ggml_tensor * root = node->view_src ? node->view_src : node; + if (!ggml_cuda_affine_unary_range(root, ctx.device, begin, end) || (begin < we && wb < end)) { return false; } + } + if (!ggml_cuda_affine_unary_view_op(node->op) && node->op != GGML_OP_NONE && ((begin < ie && ib < end) || (begin < we && wb < end))) { return false; } + for (const ggml_tensor * src : node->src) { + if (src && (!ggml_cuda_affine_unary_range(src, ctx.device, begin, end) || (begin < we && wb < end))) { return false; } + if (normalization && src) { + const ggml_tensor * root = src->view_src ? src->view_src : src; + if (!ggml_cuda_affine_unary_range(root, ctx.device, begin, end) || (begin < we && wb < end)) { return false; } + } + } + } + } + return true; +} + struct ggml_cuda_pre_mul_norm_match { ggml_tensor * repeat = nullptr; ggml_tensor * product = nullptr; @@ -5251,12 +5313,13 @@ struct ggml_cuda_pre_mul_norm_match { ggml_tensor * weighted = nullptr; int norm_i = -1; int count = 0; + int product_i = -1; }; static ggml_cuda_pre_mul_norm_match ggml_cuda_match_pre_mul_norm(ggml_backend_cuda_context & ctx, - ggml_cgraph * graph, int i, bool allocated = true) { + ggml_cgraph * graph, int i, bool allocated = true, int consumer = -1) { ggml_tensor * repeat = graph->nodes[i]->op == GGML_OP_REPEAT ? graph->nodes[i] : nullptr; - const int m = i + int(repeat != nullptr); + const int m = consumer < 0 ? i + int(repeat != nullptr) : consumer; if (m >= graph->n_nodes) { return {}; } ggml_tensor * product = graph->nodes[m]; if (product->op != GGML_OP_MUL || !product->src[0] || !product->src[1] || product->type != GGML_TYPE_F32 || @@ -5279,14 +5342,19 @@ static ggml_cuda_pre_mul_norm_match ggml_cuda_match_pre_mul_norm(ggml_backend_cu !ggml_are_same_shape(norm, norm->src[0]) || norm->ne[0] > INT_MAX - 1023 || norm->ne[2] > 65535 || norm->ne[3] > 65535) { return {}; } ggml_tensor * weighted = nullptr; int end = n + 1; - while (end < graph->n_nodes && ggml_cuda_affine_unary_view_op(graph->nodes[end]->op)) { - if (!ggml_cuda_affine_unary_prior_view(graph, i, graph->nodes[end], true)) { return {}; } + while (end < graph->n_nodes && (ggml_cuda_affine_unary_view_op(graph->nodes[end]->op) || graph->nodes[end]->op == GGML_OP_NONE)) { + const ggml_tensor * meta = graph->nodes[end]; + if (meta->op == GGML_OP_NONE) { + if (meta->flags & GGML_TENSOR_FLAG_COMPUTE) { return {}; } + } else if (!ggml_cuda_affine_unary_prior_view(graph, m, meta, true)) { return {}; } ++end; } if (end < graph->n_nodes && graph->nodes[end]->op == GGML_OP_MUL && (graph->nodes[end]->src[0] == norm) != (graph->nodes[end]->src[1] == norm)) { weighted = graph->nodes[end++]; if (!ggml_are_same_shape(norm, weighted) || !ggml_is_contiguous(weighted)) { return {}; } + } else { + end = n + 1; } // Leave longer normalization fusions to the existing dispatch. if (end < graph->n_nodes && (graph->nodes[end]->op == GGML_OP_ROPE || graph->nodes[end]->op == GGML_OP_ADD || graph->nodes[end]->op == GGML_OP_SCALE)) { return {}; } @@ -5309,10 +5377,13 @@ static ggml_cuda_pre_mul_norm_match ggml_cuda_match_pre_mul_norm(ggml_backend_cu } std::vector indices; std::vector ops; - for (int j = i; j < end; ++j) { indices.push_back(j); ops.push_back(graph->nodes[j]->op); } + if (repeat) { indices.push_back(i); ops.push_back(GGML_OP_REPEAT); } + for (int j = m; j < end; ++j) { + if (graph->nodes[j]->op != GGML_OP_NONE) { indices.push_back(j); ops.push_back(graph->nodes[j]->op); } + } if (!ggml_can_fuse_subgraph_ext(graph, indices.data(), indices.size(), ops.data(), indices.data(), indices.size())) { return {}; } if (allocated) { - for (int j = i; j < end; ++j) { + for (int j : indices) { uintptr_t begin, stop; if (!ggml_cuda_affine_unary_range(graph->nodes[j], ctx.device, begin, stop) || begin % sizeof(float)) { return {}; } } @@ -5332,7 +5403,27 @@ static ggml_cuda_pre_mul_norm_match ggml_cuda_match_pre_mul_norm(ggml_backend_cu } } } - return {repeat, product, norm, weighted, n, end - i}; + return {repeat, product, norm, weighted, n, end - (consumer < 0 ? i : m), m}; +} + +static ggml_cuda_pre_mul_norm_match ggml_cuda_match_deferred_repeat_norm(ggml_backend_cuda_context & ctx, + ggml_cgraph * graph, int i, bool allocated = true) { + const ggml_tensor * repeat = graph->nodes[i]; + if (repeat->op != GGML_OP_REPEAT || !repeat->src[0]) { return {}; } + bool compute; + const int m = ggml_cuda_repeat_first_consumer(graph, i, compute); + if (m <= i + 1 || m == graph->n_nodes) { return {}; } + const auto match = ggml_cuda_match_pre_mul_norm(ctx, graph, i, allocated, m); + if (!match.count) { return {}; } + // Do not displace an earlier producer fusion that consumes the product. + const ggml_tensor * prior = graph->nodes[m - 1]; + if (!ggml_cuda_affine_unary_view_op(prior->op) && prior->op != GGML_OP_SCALE && prior->op != GGML_OP_CONT && prior->op != GGML_OP_NONE) { + if (prior->op != GGML_OP_UNARY || ggml_are_same_shape(prior, repeat)) { return {}; } + } + ggml_cuda_moe_weighted_reduction_match reduction; + if (ggml_cuda_match_moe_weighted_reduction(graph, m, reduction) || + !ggml_cuda_repeat_interval_ok(ctx, graph, i, m, {match.repeat, match.product, match.norm, match.weighted}, allocated, true)) { return {}; } + return match; } struct ggml_cuda_repeat_mul_add_match { @@ -5391,15 +5482,8 @@ static ggml_cuda_repeat_mul_add_match ggml_cuda_match_repeat_mul_add(ggml_backen static ggml_cuda_repeat_mul_add_match ggml_cuda_match_deferred_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { ggml_tensor * repeat = graph->nodes[i]; if (repeat->op != GGML_OP_REPEAT || !repeat->src[0]) { return {}; } - int m = i + 1; - bool compute = false; - for (; m < graph->n_nodes; ++m) { - const ggml_tensor * node = graph->nodes[m]; - bool used = node->view_src == repeat; - for (const ggml_tensor * src : node->src) { used |= src && (src == repeat || src->view_src == repeat); } - if (used) { break; } - compute |= !ggml_cuda_affine_unary_view_op(node->op); - } + bool compute; + const int m = ggml_cuda_repeat_first_consumer(graph, i, compute); if (!compute || m == graph->n_nodes) { return {}; } const auto match = ggml_cuda_match_repeat_mul_add(ctx, graph, i, allocated, m); if (!match.count) { return {}; } @@ -5408,33 +5492,7 @@ static ggml_cuda_repeat_mul_add_match ggml_cuda_match_deferred_repeat_mul_add(gg if (!ggml_cuda_affine_unary_view_op(prior->op) && prior->op != GGML_OP_SCALE && prior->op != GGML_OP_CONT) { return {}; } ggml_cuda_moe_weighted_reduction_match reduction; if (ggml_cuda_match_moe_weighted_reduction(graph, m, reduction)) { return {}; } - for (int j = i + 1; j < m; ++j) { - const ggml_tensor * node = graph->nodes[j]; - switch (node->op) { - case GGML_OP_VIEW: case GGML_OP_RESHAPE: case GGML_OP_PERMUTE: case GGML_OP_TRANSPOSE: - case GGML_OP_MUL_MAT: case GGML_OP_SCALE: case GGML_OP_UNARY: case GGML_OP_MUL: - case GGML_OP_SUM_ROWS: case GGML_OP_SUB: case GGML_OP_CONT: break; - default: return {}; - } - } - if (allocated) { - const ggml_tensor * input = repeat->src[0]->view_src ? repeat->src[0]->view_src : repeat->src[0]; - uintptr_t ib, ie, wb, we; - if (!ggml_cuda_affine_unary_range(input, ctx.device, ib, ie) || !ggml_cuda_affine_unary_range(repeat, ctx.device, wb, we)) { return {}; } - for (const ggml_tensor * write : {repeat, match.binary.mul, match.binary.add}) { - uintptr_t begin, end; - if (!ggml_cuda_affine_unary_range(write, ctx.device, begin, end) || (begin < ie && ib < end)) { return {}; } - } - for (int j = i + 1; j < m; ++j) { - const ggml_tensor * node = graph->nodes[j]; - uintptr_t begin, end; - if (!ggml_cuda_affine_unary_range(node, ctx.device, begin, end)) { return {}; } - if (!ggml_cuda_affine_unary_view_op(node->op) && ((begin < ie && ib < end) || (begin < we && wb < end))) { return {}; } - for (const ggml_tensor * src : node->src) { - if (src && (!ggml_cuda_affine_unary_range(src, ctx.device, begin, end) || (begin < we && wb < end))) { return {}; } - } - } - } + if (!ggml_cuda_repeat_interval_ok(ctx, graph, i, m, {repeat, match.binary.mul, match.binary.add}, allocated)) { return {}; } return match; } @@ -6618,6 +6676,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud }; std::unordered_map repeat_emits; + std::unordered_map repeat_norm_emits; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -6662,12 +6721,25 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (!disable_reuse && cuda_ctx->curr_stream_no == 0 && stream_ctx.concurrent_events.empty()) { if (node->op == GGML_OP_REPEAT) { + const auto norm = ggml_cuda_match_deferred_repeat_norm(*cuda_ctx, cgraph, i); + if (norm.count && (norm_emits.empty() || !norm_emits[norm.norm_i].norm) && + (q8_emits.empty() || !q8_emits[norm.norm_i].norm) && (mmq_emits.empty() || !mmq_emits[norm.norm_i].norm)) { + repeat_norm_emits.emplace(norm.product, norm); + continue; + } const auto deferred = ggml_cuda_match_deferred_repeat_mul_add(*cuda_ctx, cgraph, i); if (deferred.count) { repeat_emits.emplace(deferred.binary.mul, deferred); continue; } } + const auto norm_emit = repeat_norm_emits.find(node); + if (norm_emit != repeat_norm_emits.end()) { + const auto & match = norm_emit->second; + ggml_cuda_op_rms_norm_pre_mul(*cuda_ctx, match.product, match.norm, match.weighted, match.repeat); + i += match.count - 1; + continue; + } const auto emit = repeat_emits.find(node); if (emit != repeat_emits.end()) { const auto & match = emit->second; @@ -7107,6 +7179,15 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, hc_up.mm->src[1], hc_up.dst); } + const auto repeat_norm = ggml_cuda_match_deferred_repeat_norm(*cuda_ctx, cgraph, i, false); + if (repeat_norm.count) { + ggml_tensor * last = cgraph->nodes[repeat_norm.product_i + repeat_norm.count - 1]; + ggml_tensor * read = repeat_norm.repeat->src[0]; + params->add_alloc_dep(params->user_data, read, last); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, last); } + params->add_alloc_dep(params->user_data, repeat_norm.repeat, last); + } + const auto pre = ggml_cuda_match_pre_mul_norm(*cuda_ctx, cgraph, i, false); if (pre.count) { const int last = i + pre.count - 1; From e7cff2e03e058030dcff5a7a8f531326e20fc70d Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 22:13:11 -0400 Subject: [PATCH 18/27] cuda : fuse repeat into ADD Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 122 +++++++++++++++++++++++++++++++- ggml/src/ggml-cuda/unary.cu | 48 +++++++++++++ ggml/src/ggml-cuda/unary.cuh | 2 + tests/test-backend-ops.cpp | 31 ++++++-- 4 files changed, 195 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 4f6d3fa9797..b5fbb200163 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5259,7 +5259,7 @@ static int ggml_cuda_repeat_first_consumer(ggml_cgraph * graph, int i, bool & co } static bool ggml_cuda_repeat_interval_ok(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, int m, - std::initializer_list writes, bool allocated, bool normalization = false) { + std::initializer_list writes, bool allocated, bool normalization = false, const ggml_tensor * paired = nullptr) { const ggml_tensor * repeat = graph->nodes[i]; for (int j = i + 1; j < m; ++j) { const ggml_tensor * node = graph->nodes[j]; @@ -5267,6 +5267,9 @@ static bool ggml_cuda_repeat_interval_ok(ggml_backend_cuda_context & ctx, ggml_c case GGML_OP_VIEW: case GGML_OP_RESHAPE: case GGML_OP_PERMUTE: case GGML_OP_TRANSPOSE: case GGML_OP_MUL_MAT: case GGML_OP_SCALE: case GGML_OP_UNARY: case GGML_OP_MUL: case GGML_OP_SUM_ROWS: case GGML_OP_SUB: case GGML_OP_CONT: break; + case GGML_OP_REPEAT: + if (node != paired) { return false; } + break; case GGML_OP_NONE: if (!normalization || (node->flags & GGML_TENSOR_FLAG_COMPUTE)) { return false; } break; @@ -5496,6 +5499,91 @@ static ggml_cuda_repeat_mul_add_match ggml_cuda_match_deferred_repeat_mul_add(gg return match; } +struct ggml_cuda_repeat_add_match { + ggml_tensor * repeat[2]{}; + ggml_tensor * add = nullptr; +}; + +static ggml_cuda_repeat_add_match ggml_cuda_match_repeat_add(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + ggml_tensor * first = graph->nodes[i]; + if (first->op != GGML_OP_REPEAT || !first->src[0]) { return {}; } + bool compute; + const int a = ggml_cuda_repeat_first_consumer(graph, i, compute); + if (a >= graph->n_nodes) { return {}; } + ggml_tensor * add = graph->nodes[a]; + if (add->op != GGML_OP_ADD || !ggml_cuda_affine_unary_binary_type(add) || !ggml_is_contiguous(add) || + ggml_nelements(add) <= 0 || ggml_nelements(add) > INT_MAX - CUDA_NEG_BLOCK_SIZE) { return {}; } + if (a + 1 < graph->n_nodes && ggml_can_fuse(graph, a, {GGML_OP_ADD, GGML_OP_ADD}) && + graph->nodes[a + 1]->src[0] == add && ggml_are_same_layout(add->src[1], graph->nodes[a + 1]->src[1])) { return {}; } + int prior = a - 1; + while (prior > i && (ggml_cuda_affine_unary_view_op(graph->nodes[prior]->op) || graph->nodes[prior]->op == GGML_OP_NONE)) { --prior; } + if (prior > i) { + // Keep earlier producer and ADD-chain fusions. + switch (graph->nodes[prior]->op) { + case GGML_OP_MUL: case GGML_OP_ADD: case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: + case GGML_OP_RMS_NORM: case GGML_OP_SSM_CONV: return {}; + default: break; + } + } + ggml_cuda_repeat_add_match match; + match.add = add; + int indices[3], count = 0; + ggml_op ops[3]; + std::vector writes = {add}; + const ggml_tensor * reads[2]; + for (int j = 0; j < 2; ++j) { + ggml_tensor * input = add->src[j]; + int r = i; + while (r < a && graph->nodes[r] != input) { ++r; } + if (input->op == GGML_OP_REPEAT && r < a) { + bool gap; + if (!input->src[0] || (input->type != GGML_TYPE_F32 && input->type != GGML_TYPE_F16) || + input->type != input->src[0]->type || !ggml_is_contiguous(input) || !ggml_are_same_shape(input, add) || ggml_cuda_repeat_first_consumer(graph, r, gap) != a) { return {}; } + match.repeat[j] = input; + if (j == 0 || input != match.repeat[0]) { + indices[count] = r; ops[count++] = GGML_OP_REPEAT; + writes.push_back(input); + } + } + reads[j] = match.repeat[j] ? input->src[0] : input; + const ggml_tensor * read = reads[j]; + size_t bytes; + if (!ggml_cuda_affine_unary_bytes(read, bytes)) { return {}; } + const size_t size = ggml_type_size(read->type); + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (read->ne[d] > INT_MAX || read->nb[d] % size || read->nb[d]/size > UINT32_MAX || add->ne[d] % read->ne[d]) { return {}; } + } + } + if (first != match.repeat[0] && first != match.repeat[1]) { return {}; } + if (count == 2 && indices[0] > indices[1]) { std::swap(indices[0], indices[1]); } + indices[count] = a; ops[count++] = GGML_OP_ADD; + if (!ggml_can_fuse_subgraph_ext(graph, indices, count, ops, indices, count)) { return {}; } + for (int j = 0; j < 2; ++j) { + if (!match.repeat[j] || (j == 1 && match.repeat[1] == match.repeat[0])) { continue; } + int r = i; + while (r < a && graph->nodes[r] != match.repeat[j]) { ++r; } + if (!ggml_cuda_repeat_interval_ok(ctx, graph, r, a, {add, match.repeat[0], match.repeat[1]}, allocated, true, match.repeat[1 - j])) { return {}; } + } + for (const ggml_tensor * write : writes) { + if (!(write->flags & GGML_TENSOR_FLAG_COMPUTE)) { return {}; } + for (const ggml_tensor * read : reads) { if (read == write || read->view_src == write) { return {}; } } + } + if (allocated) { + std::vector wb(writes.size()), we(writes.size()); + for (size_t j = 0; j < writes.size(); ++j) { + if (!ggml_cuda_affine_unary_range(writes[j], ctx.device, wb[j], we[j]) || wb[j] % ggml_type_size(writes[j]->type)) { return {}; } + for (size_t k = 0; k < j; ++k) { if (wb[j] < we[k] && wb[k] < we[j]) { return {}; } } + for (const ggml_tensor * read : reads) { + const ggml_tensor * root = read->view_src ? read->view_src : read; + uintptr_t rb, re; + if (!ggml_cuda_affine_unary_range(read, ctx.device, rb, re) || rb % ggml_type_size(read->type) || + !ggml_cuda_affine_unary_range(root, ctx.device, rb, re) || (wb[j] < re && rb < we[j])) { return {}; } + } + } + } + return match; +} + struct ggml_cuda_ordered_mul_add_match { ggml_tensor * products[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; ggml_tensor * sums[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; @@ -6677,6 +6765,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud std::unordered_map repeat_emits; std::unordered_map repeat_norm_emits; + std::unordered_map repeat_add_emits; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; if (is_concurrent_event_active) { @@ -6720,6 +6809,13 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } if (!disable_reuse && cuda_ctx->curr_stream_no == 0 && stream_ctx.concurrent_events.empty()) { + const auto add_emit = repeat_add_emits.find(node); + if (add_emit != repeat_add_emits.end()) { + if (node->op == GGML_OP_REPEAT) { continue; } + const auto & match = add_emit->second; + ggml_cuda_op_repeat_add(*cuda_ctx, match.add, match.repeat[0], match.repeat[1]); + continue; + } if (node->op == GGML_OP_REPEAT) { const auto norm = ggml_cuda_match_deferred_repeat_norm(*cuda_ctx, cgraph, i); if (norm.count && (norm_emits.empty() || !norm_emits[norm.norm_i].norm) && @@ -6732,6 +6828,14 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud repeat_emits.emplace(deferred.binary.mul, deferred); continue; } + const auto repeated_add = ggml_cuda_match_repeat_add(*cuda_ctx, cgraph, i); + if (repeated_add.add) { + repeat_add_emits.emplace(repeated_add.add, repeated_add); + for (ggml_tensor * repeat : repeated_add.repeat) { + if (repeat && repeat != node) { repeat_add_emits.emplace(repeat, repeated_add); } + } + continue; + } } const auto norm_emit = repeat_norm_emits.find(node); if (norm_emit != repeat_norm_emits.end()) { @@ -7213,6 +7317,22 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, deferred.binary.mul->src[1], deferred.binary.add); } + const auto repeated_add = ggml_cuda_match_repeat_add(*cuda_ctx, cgraph, i, false); + if (repeated_add.add) { + for (ggml_tensor * repeat : repeated_add.repeat) { + if (!repeat) { continue; } + ggml_tensor * read = repeat->src[0]; + params->add_alloc_dep(params->user_data, read, repeated_add.add); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, repeated_add.add); } + params->add_alloc_dep(params->user_data, repeat, repeated_add.add); + } + for (ggml_tensor * read : repeated_add.add->src) { + if (!read) { continue; } + params->add_alloc_dep(params->user_data, read, repeated_add.add); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, repeated_add.add); } + } + } + const auto repeated = ggml_cuda_match_repeat_mul_add(*cuda_ctx, cgraph, i, false); if (repeated.count) { for (ggml_tensor * read : {repeated.repeat->src[0], repeated.repeat, repeated.binary.mul->src[0], repeated.binary.mul->src[1], repeated.binary.add->src[repeated.binary.add->src[0] == repeated.binary.mul ? 1 : 0]}) { diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 7d3907d0633..5c18d302b7b 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -494,6 +494,54 @@ void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * ggml_cuda_kernel_launch(repeat_mul_add_kernel, launch, binary, mul->src[0] == repeat ? 0 : 1, add->src[0] == mul); } +static __global__ void repeat_add_kernel(mul_add_args binary) { + const affine_unary_args & args = binary.affine; + ggml_cuda_pdl_lc(); + const int i = blockDim.x*blockIdx.x + threadIdx.x; + if (i >= args.count) { return; } + uint32_t remaining = i; + uint32_t pos[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + const uint2 part = fast_div_modulo(remaining, binary.ne[d]); + remaining = part.x; + pos[d] = part.y; + } + ggml_cuda_pdl_sync(); + float lhs = mul_add_read(args.input[0], binary.repeat[0], pos, binary.repeat_mask[0]); + float rhs = mul_add_read(args.input[1], binary.repeat[1], pos, binary.repeat_mask[1]); + if (args.output[1]) { lhs = affine_unary_store(args.output[1], args.output_type[1], i, lhs); } + if (args.output[2]) { rhs = affine_unary_store(args.output[2], args.output_type[2], i, rhs); } + affine_unary_store(args.output[0], args.output_type[0], i, lhs + rhs); +} + +void ggml_cuda_op_repeat_add(ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * repeat0, ggml_tensor * repeat1) { + mul_add_args binary{}; + affine_unary_args & args = binary.affine; + const ggml_tensor * repeats[] = {repeat0, repeat1}; + for (int j = 0; j < 2; ++j) { + const ggml_tensor * read = repeats[j] ? repeats[j]->src[0] : add->src[j]; + args.input[j].data = read->data; + args.input[j].type = read->type; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + args.input[j].ne[d] = read->ne[d]; + args.input[j].nb[d] = read->ne[d] == 1 ? 0 : read->nb[d]; + binary.repeat[j][d] = init_fastdiv_values(read->ne[d]); + if (read->ne[d] != 1 && read->ne[d] != add->ne[d]) { binary.repeat_mask[j] |= 1u << d; } + } + if (repeats[j] && (j == 0 || repeat1 != repeat0)) { + args.output[j + 1] = repeats[j]->data; + args.output_type[j + 1] = repeats[j]->type; + } + } + args.output[0] = add->data; + args.output_type[0] = add->type; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { binary.ne[d] = init_fastdiv_values(add->ne[d]); } + args.count = ggml_nelements(add); + const int blocks = (args.count + CUDA_NEG_BLOCK_SIZE - 1)/CUDA_NEG_BLOCK_SIZE; + const auto launch = ggml_cuda_kernel_launch_params(blocks, CUDA_NEG_BLOCK_SIZE, 0, ctx.stream()); + ggml_cuda_kernel_launch(repeat_add_kernel, launch, binary); +} + struct ordered_mul_add_args { const char * input[2]; size_t row_stride[2]; diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 635d774ab51..5a854d20618 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -96,6 +96,8 @@ void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary void ggml_cuda_op_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * add); void ggml_cuda_op_repeat_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * repeat, ggml_tensor * mul, ggml_tensor * add); +void ggml_cuda_op_repeat_add(ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * repeat0, ggml_tensor * repeat1); + constexpr int GGML_CUDA_ORDERED_MUL_ADD_MAX = 128; void ggml_cuda_op_ordered_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * const * products, ggml_tensor * const * sums, int count); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 25cfeae9a29..f34f07aebd4 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3302,10 +3302,11 @@ struct test_bin_bcast : public test_case { const bool coefficient_gap; std::vector stages; - bool run_whole_graph() override { return nf > 1 || add_after; } + bool run_whole_graph() override { return nf > 1 || add_after || repeat_input >= 0; } std::vector fusion_test_nodes() override { return stages; } std::string op_desc(ggml_tensor * t) override { + if (op == ggml_add && repeat_input >= 0) { return coefficient_gap ? "DEFERRED_REPEAT_ADD" : "REPEAT_ADD"; } return coefficient_gap ? "DEFERRED_REPEAT_MUL_ADD" : repeat_input >= 0 ? "REPEAT_MUL_ADD" : add_after ? "MUL_ADD" : test_case::op_desc(t); } @@ -3337,7 +3338,7 @@ struct test_bin_bcast : public test_case { stages.clear(); std::array expanded = {ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]}; - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, repeat_input == 0 ? ne.data() : expanded.data()); + ggml_tensor * a = ggml_new_tensor(ctx, type, 4, (repeat_input == 0 || repeat_input == 2) ? ne.data() : expanded.data()); ggml_set_name(a, "a"); const ggml_type tb = type_b == GGML_TYPE_COUNT ? type : type_b; @@ -3358,16 +3359,16 @@ struct test_bin_bcast : public test_case { } if (repeat_input >= 0) { - GGML_ASSERT(op == ggml_mul && nf == 1 && add_after && !src_overlap && repeat_input <= 1); - ggml_tensor * repeated = ggml_repeat_4d(ctx, repeat_input == 0 ? a : b[0], expanded[0], expanded[1], expanded[2], expanded[3]); + GGML_ASSERT(nf == 1 && !src_overlap && ((op == ggml_mul && add_after && repeat_input <= 1) || (op == ggml_add && !add_after && repeat_input <= 2))); + ggml_tensor * repeated = ggml_repeat_4d(ctx, repeat_input == 1 ? b[0] : a, expanded[0], expanded[1], expanded[2], expanded[3]); ggml_set_output(repeated); stages.push_back(repeated); - if (repeat_input == 0) a = repeated; + if (repeat_input != 1) a = repeated; else b[0] = repeated; } if (coefficient_gap) { - GGML_ASSERT(repeat_input == 0 && tb == GGML_TYPE_F32 && !perm1 && !src_overlap); + GGML_ASSERT((repeat_input == 0 || repeat_input == 2) && tb == GGML_TYPE_F32 && !perm1 && !src_overlap); b[0] = ggml_scale(ctx, b[0], 0.25f); ggml_set_output(b[0]); stages.push_back(b[0]); b[0] = ggml_sigmoid(ctx, b[0]); @@ -3376,8 +3377,13 @@ struct test_bin_bcast : public test_case { ggml_set_output(b[0]); stages.push_back(b[0]); } + if (repeat_input == 2) { + b[0] = ggml_repeat_4d(ctx, b[0], expanded[0], expanded[1], expanded[2], expanded[3]); + ggml_set_output(b[0]); stages.push_back(b[0]); + } + // The backward pass supports broadcasting only for GGML_ADD: - const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; + const bool grad_supported = op == ggml_add && repeat_input < 0 && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; if (grad_supported) { ggml_set_param(a); ggml_set_param(b[0]); @@ -3395,6 +3401,8 @@ struct test_bin_bcast : public test_case { out = op(ctx, out, b[i]); } + if (op == ggml_add && repeat_input >= 0) { ggml_set_output(out); stages.push_back(out); } + if (add_after) { GGML_ASSERT(op == ggml_mul && nf == 1); ggml_set_output(out); @@ -10112,6 +10120,15 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { + for (int repeated : {0, 1, 2}) { + test_cases.emplace_back(new test_bin_bcast(ggml_add, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, false, false, GGML_TYPE_COUNT, false, false, repeated)); + } + test_cases.emplace_back(new test_bin_bcast(ggml_add, type, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, false, false, GGML_TYPE_F32, false, false, 2, true)); + } + + test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_BF16, {7, 2, 3, 2}, {3, 2, 2, 2}, 1, false, false, GGML_TYPE_F32, false, false, 1)); + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { for (bool permute : {false, true}) { for (bool swap : {false, true}) { From 0ea1373e74b04ab5547d4c0619324786e58eef0a Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 22:53:47 -0400 Subject: [PATCH 19/27] cuda : read repeated matrix banks from compact inputs Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 80 +++++++++++++++++++++++++++++++++ tests/test-backend-ops.cpp | 30 ++++++++++--- 2 files changed, 104 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index b5fbb200163..6fe927ed0b9 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5584,6 +5584,64 @@ static ggml_cuda_repeat_add_match ggml_cuda_match_repeat_add(ggml_backend_cuda_c return match; } +struct ggml_cuda_repeat_mm_match { + ggml_tensor * repeat = nullptr; + ggml_tensor * mm = nullptr; +}; + +static ggml_cuda_repeat_mm_match ggml_cuda_match_repeat_mm(ggml_backend_cuda_context & ctx, ggml_cgraph * graph, int i, bool allocated = true) { + const int indices[] = {i, i + 1}; + const ggml_op ops[] = {GGML_OP_REPEAT, GGML_OP_MUL_MAT}; + if (!ggml_can_fuse_subgraph_ext(graph, indices, 2, ops, indices + 1, 1)) { return {}; } + ggml_tensor * repeat = graph->nodes[i]; + ggml_tensor * mm = graph->nodes[i + 1]; + const ggml_tensor * input = repeat->src[0]; + if (!input || repeat->type != GGML_TYPE_F32 || input->type != GGML_TYPE_F32 || mm->type != GGML_TYPE_F32 || + mm->src[1] != repeat || !mm->src[0] || !ggml_is_contiguous(repeat) || !ggml_is_contiguous(mm) || + input->nb[0] != sizeof(float) || input->nb[1] % (2*sizeof(float)) || ggml_cuda_op_mul_mat_use_fwht(mm)) { return {}; } + bool expanded = false; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (input->ne[d] != repeat->ne[d]) { + if (d < 2 || input->ne[d] != 1) { return {}; } + expanded = true; + } + if (input->ne[d] <= 0 || repeat->ne[d] > INT_MAX || input->nb[d] % sizeof(float) || input->nb[d]/sizeof(float) > INT_MAX) { return {}; } + } + size_t bytes; + if (!expanded || !ggml_cuda_affine_unary_bytes(input, bytes) || bytes/sizeof(float) > INT_MAX) { return {}; } + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + if (input->ne[d] > 1 && input->nb[d] % (2*sizeof(float))) { return {}; } + } + const auto kernel = ggml_cuda_mul_mat_kernel(mm->src[0], repeat, mm, ctx.device); + if (kernel != GGML_CUDA_MM_MMVF && kernel != GGML_CUDA_MM_MMF) { return {}; } + int next = i + 2; + while (next < graph->n_nodes && (ggml_cuda_affine_unary_view_op(graph->nodes[next]->op) || graph->nodes[next]->op == GGML_OP_NONE)) { ++next; } + if (next < graph->n_nodes) { + // Keep existing matrix producer fusions. + switch (graph->nodes[next]->op) { + case GGML_OP_ADD: case GGML_OP_MUL: case GGML_OP_SCALE: case GGML_OP_UNARY: + case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: case GGML_OP_GLU: case GGML_OP_DSV4_HC_PRE: return {}; + default: break; + } + } + if (allocated) { + const ggml_tensor * reads[] = {input, mm->src[0]}; + uintptr_t wb[2], we[2]; + for (int j = 0; j < 2; ++j) { + const ggml_tensor * write = j ? mm : repeat; + if (!ggml_cuda_affine_unary_range(write, ctx.device, wb[j], we[j]) || wb[j] % sizeof(float)) { return {}; } + for (const ggml_tensor * read : reads) { + const ggml_tensor * root = read->view_src ? read->view_src : read; + uintptr_t rb, re; + if (!ggml_cuda_affine_unary_range(read, ctx.device, rb, re) || rb % (2*ggml_type_size(read->type)) || + !ggml_cuda_affine_unary_range(root, ctx.device, rb, re) || (wb[j] < re && rb < we[j])) { return {}; } + } + } + if (wb[0] < we[1] && wb[1] < we[0]) { return {}; } + } + return {repeat, mm}; +} + struct ggml_cuda_ordered_mul_add_match { ggml_tensor * products[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; ggml_tensor * sums[GGML_CUDA_ORDERED_MUL_ADD_MAX]{}; @@ -5704,6 +5762,21 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if (cuda_ctx->curr_stream_no == 0 && cuda_ctx->stream_context().concurrent_events.empty()) { + const auto match = ggml_cuda_match_repeat_mm(*cuda_ctx, cgraph, i); + if (match.mm) { + const ggml_tensor * input = match.repeat->src[0]; + ggml_tensor reader = *match.repeat; + reader.data = input->data; + reader.buffer = input->buffer; + reader.view_src = input->view_src; + reader.view_offs = input->view_offs; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { reader.nb[d] = input->ne[d] == 1 ? 0 : input->nb[d]; } + ggml_cuda_mul_mat(*cuda_ctx, match.mm->src[0], &reader, match.mm); + return 1; + } + } + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache if (node->op == GGML_OP_GATED_DELTA_NET) { ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; @@ -7317,6 +7390,13 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph params->add_alloc_dep(params->user_data, deferred.binary.mul->src[1], deferred.binary.add); } + const auto repeat_mm = ggml_cuda_match_repeat_mm(*cuda_ctx, cgraph, i, false); + if (repeat_mm.mm) { + ggml_tensor * read = repeat_mm.repeat->src[0]; + params->add_alloc_dep(params->user_data, read, repeat_mm.mm); + if (read->view_src) { params->add_alloc_dep(params->user_data, read->view_src, repeat_mm.mm); } + } + const auto repeated_add = ggml_cuda_match_repeat_add(*cuda_ctx, cgraph, i, false); if (repeated_add.add) { for (ggml_tensor * repeat : repeated_add.repeat) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f34f07aebd4..2fd9c3dacdd 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5191,8 +5191,12 @@ struct test_mul_mat : public test_case { const int64_t m_v; // rows of a in memory, the batches of a are strided for m_v > m, no view for m_v == 0 const int64_t pad; // bytes after the m_v rows of each batch of a, so nb[2] of a is not a multiple of nb[1] + const bool repeat_rhs; + std::string vars() override { - return VARS_TO_STR13(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, m_v, pad); + std::string s = VARS_TO_STR13(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, m_v, pad); + if (repeat_rhs) { s += "," + VAR_TO_STR(repeat_rhs); } + return s; } double max_nmse_err() override { @@ -5223,8 +5227,8 @@ struct test_mul_mat : public test_case { std::array bs = {10, 10}, std::array nr = {2, 2}, std::array per = {0, 1, 2, 3}, - int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad) {} + int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0, bool repeat_rhs = false) + : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad), repeat_rhs(repeat_rhs) {} ggml_tensor * build_graph(ggml_context * ctx) override { // C^T = A * B^T: (k, m) * (k, n) => (m, n) @@ -5232,7 +5236,15 @@ struct test_mul_mat : public test_case { ggml_tensor * b; const int npermuted = (per[0] != 0) + (per[1] != 1) + (per[2] != 2) + (per[3] != 3); - if (npermuted > 0) { + if (repeat_rhs) { + GGML_ASSERT(type_b == GGML_TYPE_F32 && npermuted == 0 && !src_overlap && k_v == 0 && m_v == 0 && pad == 0 && o == 1); + a = ggml_new_tensor_4d(ctx, type_a, k, m, bs[0], bs[1]); + b = ggml_new_tensor_2d(ctx, type_b, k, n); + ggml_set_name(a, "a"); + ggml_set_name(b, "b_compact"); + b = ggml_repeat_4d(ctx, b, k, n, bs[0]*nr[0], bs[1]*nr[1]); + ggml_set_name(b, "b_repeated"); + } else if (npermuted > 0) { GGML_ASSERT(npermuted == 2); GGML_ASSERT(k_v == 0); // not handled GGML_ASSERT(m_v == 0); // not handled @@ -5309,11 +5321,11 @@ struct test_mul_mat : public test_case { return out; } - bool run_whole_graph() override { return o > 1; } + bool run_whole_graph() override { return o > 1 || repeat_rhs; } std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); - return ggml_op_name(GGML_OP_MUL_MAT); + return repeat_rhs ? "REPEAT_MUL_MAT" : ggml_op_name(GGML_OP_MUL_MAT); } }; @@ -10448,6 +10460,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { + for (int tokens : {1, 3, 8}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, tokens, 128, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 0, 0, true)); + } + } + #if 0 // > 4GB A matrix. Too slow to be enabled by default. test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 900000, 3, 2592, {1, 1}, {1, 1})); From e9f2bfdc272d4ad477ca7104592d006c35702b04 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Tue, 6 Oct 2026 23:43:15 -0400 Subject: [PATCH 20/27] cuda : read repeated quantized matrix banks from compact inputs Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 27 +++++++++++++++++++++++++-- tests/test-backend-ops.cpp | 7 +++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6fe927ed0b9..55c9bc94f1d 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5613,7 +5613,8 @@ static ggml_cuda_repeat_mm_match ggml_cuda_match_repeat_mm(ggml_backend_cuda_con if (input->ne[d] > 1 && input->nb[d] % (2*sizeof(float))) { return {}; } } const auto kernel = ggml_cuda_mul_mat_kernel(mm->src[0], repeat, mm, ctx.device); - if (kernel != GGML_CUDA_MM_MMVF && kernel != GGML_CUDA_MM_MMF) { return {}; } + const bool quantized = kernel == GGML_CUDA_MM_MMVQ || kernel == GGML_CUDA_MM_MMQ; + if (kernel != GGML_CUDA_MM_MMVF && kernel != GGML_CUDA_MM_MMF && !quantized) { return {}; } int next = i + 2; while (next < graph->n_nodes && (ggml_cuda_affine_unary_view_op(graph->nodes[next]->op) || graph->nodes[next]->op == GGML_OP_NONE)) { ++next; } if (next < graph->n_nodes) { @@ -5633,11 +5634,33 @@ static ggml_cuda_repeat_mm_match ggml_cuda_match_repeat_mm(ggml_backend_cuda_con for (const ggml_tensor * read : reads) { const ggml_tensor * root = read->view_src ? read->view_src : read; uintptr_t rb, re; - if (!ggml_cuda_affine_unary_range(read, ctx.device, rb, re) || rb % (2*ggml_type_size(read->type)) || + if (!ggml_cuda_affine_unary_range(read, ctx.device, rb, re) || rb % (ggml_is_quantized(read->type) ? 16 : 2*ggml_type_size(read->type)) || !ggml_cuda_affine_unary_range(root, ctx.device, rb, re) || (wb[j] < re && rb < we[j])) { return {}; } } } if (wb[0] < we[1] && wb[1] < we[0]) { return {}; } + if (quantized) { + const ggml_tensor * weight = mm->src[0]; + if (ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size = ggml_backend_buffer_get_alloc_size(weight->buffer, weight); + const uintptr_t begin = (uintptr_t) weight->data; + if (size > UINTPTR_MAX - begin) { return {}; } + const uintptr_t end = begin + size; + const uintptr_t padding = begin + ggml_nbytes(weight); + const ggml_tensor * root = input->view_src ? input->view_src : input; + uintptr_t rb, re; + if (!ggml_cuda_affine_unary_range(root, ctx.device, rb, re)) { return {}; } + if (padding < end && ((padding < re && rb < end) || + (padding < we[0] && wb[0] < end) || (padding < we[1] && wb[1] < end))) { return {}; } + } + if (kernel == GGML_CUDA_MM_MMQ) { + ggml_tensor reader = *input; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { if (input->ne[d] == 1) { reader.nb[d] = 0; } } + const auto prec = ggml_cuda_mmq_get_prec_src1(weight, mm, ggml_cuda_info().devices[ctx.device].cc); + if (prec == GGML_PREC_Q8 && !ggml_cuda_is_aligned(&reader, 16)) { return {}; } + if (prec == GGML_PREC_Q4 && ggml_cuda_is_aligned(&reader, 32) != ggml_cuda_is_aligned(repeat, 32)) { return {}; } + } + } } return {repeat, mm}; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2fd9c3dacdd..c8987e2f404 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10466,6 +10466,13 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : all_types) { + if (!ggml_is_quantized(type)) { continue; } + for (int tokens : {1, 3, 16}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, tokens, 512, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 0, 0, true)); + } + } + #if 0 // > 4GB A matrix. Too slow to be enabled by default. test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 900000, 3, 2592, {1, 1}, {1, 1})); From 7ef1d05483aad1e3bf974c2cfe6b1f0ba7f41496 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Wed, 7 Oct 2026 00:42:05 -0400 Subject: [PATCH 21/27] cuda : pack broadcast activation banks once Assisted-by: Codex --- ggml/src/ggml-cuda/mmq.cu | 20 +++++++++++++++----- ggml/src/ggml-cuda/mmvq.cu | 14 +++++++++++--- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 4bc02d869d5..2ad85695f29 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -232,11 +232,20 @@ void ggml_cuda_mul_mat_q( if (!ids) { ggml_cuda_mmq_input local_input(ctx.pool()); + ggml_tensor packed = *src1; if (!input) { + // NVFP4 row scales still use the full logical bank/sample layout. + if (!use_native_fp4 || src0->type != GGML_TYPE_NVFP4) { + for (int d = 2; d < GGML_MAX_DIMS; ++d) { + if (packed.nb[d] == 0) { + packed.ne[d] = 1; + } + } + } const int J_max = ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11); - const size_t size = ne13*ne12*ne11*ne10_padded*y_block_size/y_values_per_block + + const size_t size = packed.ne[3]*packed.ne[2]*ne11*ne10_padded*y_block_size/y_values_per_block + J_max*sizeof(block_q8_1_mmq); - ggml_cuda_quantize_mmq_input_impl(ctx, src0, src1, prec_src1, size, local_input); + ggml_cuda_quantize_mmq_input_impl(ctx, src0, &packed, prec_src1, size, local_input); input = &local_input; } @@ -244,14 +253,15 @@ void ggml_cuda_mul_mat_q( const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) : ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); - const int64_t s13 = ne12*s12; + const int64_t stride_channel_y = packed.ne[2] == ne12 ? s12 : 0; + const int64_t stride_sample_y = packed.ne[3] == ne13 ? packed.ne[2]*s12 : 0; const mmq_args args = { src0_d, src0->type, (const int *) input->quantized.ptr, nullptr, nullptr, dst_d, src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? input->scale.ptr : nullptr, ne00, ne01, ne1, s01, ne11, s1, - ne02, ne12, s02, s12, s2, - ne03, ne13, s03, s13, s3, + ne02, ne12, s02, stride_channel_y, s2, + ne03, ne13, s03, stride_sample_y, s3, ne1, ne1}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1); return; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 9581632d9c2..9554f7ca890 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1507,8 +1507,16 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); + ggml_tensor packed = *src1; + if (!ids && !quantized) { + for (int d = 2; d < GGML_MAX_DIMS; ++d) { + if (packed.nb[d] == 0) { + packed.ne[d] = 1; + } + } + } if (!quantized) { - ggml_cuda_quantize_mmvq_input(ctx, src0, src1, src1_q8_1); + ggml_cuda_quantize_mmvq_input(ctx, src0, &packed, src1_q8_1); quantized = src1_q8_1.get(); } @@ -1520,8 +1528,8 @@ void ggml_cuda_mul_mat_vec_q( const int64_t s03 = src0->nb[3] / ts_src0; const int64_t s3 = dst->nb[3] / ts_dst; - const int64_t s12 = ne11*s11; - const int64_t s13 = ne12*s12; + const int64_t s12 = packed.ne[2] == ne12 ? ne11*s11 : 0; + const int64_t s13 = packed.ne[3] == ne13 ? packed.ne[2]*ne11*s11 : 0; // For MUL_MAT_ID the memory layout is different than for MUL_MAT: const int64_t ncols_dst = ids ? ne2 : ne1; From 8437ebba92304cdaf1c37f7884305a872a1ef3ea Mon Sep 17 00:00:00 2001 From: Rocco A Date: Wed, 7 Oct 2026 01:34:06 -0400 Subject: [PATCH 22/27] cuda : share banked activation images Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 35 ++++++++++++++++++++++----------- tests/test-backend-ops.cpp | 6 ++++++ 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 55c9bc94f1d..9a6ec29b222 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3086,6 +3086,15 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, return is_ok; } +static bool ggml_cuda_input_rows(const ggml_tensor * input, size_t & rows) { + rows = 1; + for (int d = 1; d < GGML_MAX_DIMS; ++d) { + if (input->ne[d] <= 0 || uint64_t(input->ne[d]) > SIZE_MAX/rows) { return false; } + rows *= input->ne[d]; + } + return true; +} + static bool ggml_cuda_can_share_mmq_input(const ggml_tensor * node, int device) { if (node->op != GGML_OP_MUL_MAT || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || node->view_src || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node)) { @@ -3095,7 +3104,7 @@ static bool ggml_cuda_can_share_mmq_input(const ggml_tensor * node, int device) const ggml_tensor * input = node->src[1]; const int cc = ggml_cuda_info().devices[device].cc; if (input->type != GGML_TYPE_F32 || input->nb[0] != sizeof(float) || input->ne[1] < 1 || - input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || + ((input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) && !ggml_is_contiguous(input)) || !ggml_is_contiguous(weight) || ggml_cuda_should_use_mmvq(weight->type, cc, input->ne[1]) || !ggml_cuda_should_use_mmq(weight->type, cc, input->ne[1], /*n_experts =*/ 0) || ggml_cuda_op_mul_mat_use_fwht(node)) { return false; @@ -3558,7 +3567,7 @@ static bool ggml_cuda_can_share_mmvq_input(const ggml_tensor * node, int device) const ggml_tensor * input = node->src[1]; const int cc = ggml_cuda_info().devices[device].cc; if (cc <= GGML_CUDA_CC_PASCAL || input->type != GGML_TYPE_F32 || input->nb[0] != sizeof(float) || - input->ne[1] < 1 || input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1 || + input->ne[1] < 1 || ((input->ne[2] != 1 || input->ne[3] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) && !ggml_is_contiguous(input)) || !ggml_is_contiguous(weight) || !ggml_cuda_should_use_mmvq(weight->type, cc, input->ne[1]) || ggml_cuda_op_mul_mat_use_fwht(node)) { return false; @@ -4146,7 +4155,7 @@ static std::vector ggml_cuda_plan_norm_q8(ggml_cgraph * const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if (prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + if (input->ne[2] != 1 || input->ne[3] != 1 || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } if (match.post && input->ne[0] % MATRIX_ROW_PADDING && match.norm->ne[0] >= input->ne[0]) { @@ -4310,7 +4319,7 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if (keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || + if (input->ne[2] != 1 || input->ne[3] != 1 || keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || norm_group_invalid(node) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } @@ -6760,9 +6769,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud const size_t blocks = GGML_PAD(size_t(input->ne[0]), MATRIX_ROW_PADDING)/QK8_1; if (blocks > SIZE_MAX/sizeof(block_q8_1)) { continue; } const size_t row = blocks*sizeof(block_q8_1); - if (uint64_t(input->ne[1]) > SIZE_MAX/row) { continue; } + size_t rows; + if (!ggml_cuda_input_rows(input, rows) || rows > SIZE_MAX/row) { continue; } mmvq_keys[i] = 0; - mmvq_sizes[i] = size_t(input->ne[1])*row; + mmvq_sizes[i] = rows*row; } } } @@ -6789,8 +6799,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud std::vector mmq_sizes; const int mmq_cc = ggml_cuda_info().devices[cuda_ctx->device].cc; const auto scale_size = [&](const ggml_tensor * node) { - return ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) == GGML_PREC_Q4 && node->src[0]->type == GGML_TYPE_NVFP4 ? - GGML_PAD((size_t(node->src[1]->ne[1]) + 128)*sizeof(float), 256) : 0; + if (ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) != GGML_PREC_Q4 || node->src[0]->type != GGML_TYPE_NVFP4) { return size_t(0); } + size_t rows; + if (!ggml_cuda_input_rows(node->src[1], rows) || rows > (SIZE_MAX - 255)/sizeof(float) - 128) { return SIZE_MAX; } + return GGML_PAD((rows + 128)*sizeof(float), 256); }; if (!disable_reuse) { mmq_keys.assign(cgraph->n_nodes, -1); @@ -6807,9 +6819,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud const size_t row = blocks*sizeof(block_q8_1_mmq); // The kernel reads a full tile, including columns past the input. const size_t tail = 128*sizeof(block_q8_1_mmq); - if (uint64_t(input->ne[1]) > (SIZE_MAX - tail - 255)/row || - uint64_t(input->ne[1]) > (SIZE_MAX - 255)/sizeof(float) - 128) { continue; } - const size_t quantized = size_t(input->ne[1])*row + tail; + size_t rows; + if (!ggml_cuda_input_rows(input, rows) || rows > (SIZE_MAX - tail - 255)/row || + rows > (SIZE_MAX - 255)/sizeof(float) - 128) { continue; } + const size_t quantized = rows*row + tail; if (scale_size(node) > SIZE_MAX - GGML_PAD(quantized, 256)) { continue; } mmq_keys[i] = ggml_cuda_mmq_get_prec_src1(node->src[0], node, mmq_cc) == GGML_PREC_Q4 ? 3 + (node->src[0]->type == GGML_TYPE_NVFP4) : int(mmq_get_q8_1_ds_layout(node->src[0]->type)); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c8987e2f404..bac536ef1d6 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10460,6 +10460,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); + for (ggml_type type : {GGML_TYPE_Q4_K, GGML_TYPE_Q4_0, GGML_TYPE_MXFP4, GGML_TYPE_NVFP4}) { + for (int tokens : {1, 17}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 65, tokens, 512, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 2)); + } + } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16}) { for (int tokens : {1, 3, 8}) { test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, tokens, 128, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 0, 0, true)); From b75964faa1b178c1fb1412699f98143488db2e54 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Wed, 7 Oct 2026 02:38:52 -0400 Subject: [PATCH 23/27] cuda : emit banked RMS activation images Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 6 +- ggml/src/ggml-cuda/norm.cu | 123 +++++++++++++++++++++++++------- ggml/src/ggml-cuda/norm.cuh | 2 +- tests/test-backend-ops.cpp | 18 ++++- 4 files changed, 115 insertions(+), 34 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 9a6ec29b222..07b31f07f68 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4155,7 +4155,7 @@ static std::vector ggml_cuda_plan_norm_q8(ggml_cgraph * const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if (input->ne[2] != 1 || input->ne[3] != 1 || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + if ((match.post && (input->ne[2] != 1 || input->ne[3] != 1)) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } if (match.post && input->ne[0] % MATRIX_ROW_PADDING && match.norm->ne[0] >= input->ne[0]) { @@ -4319,7 +4319,7 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if (input->ne[2] != 1 || input->ne[3] != 1 || keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || + if ((match.post && (input->ne[2] != 1 || input->ne[3] != 1)) || keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || norm_group_invalid(node) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } @@ -7130,7 +7130,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud const ggml_tensor * input = cgraph->nodes[reader]->src[1]; mmq_inputs[g].quantized.alloc(mmq_reuse.groups[g].size); ggml_cuda_op_rms_norm_mmq(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, - mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); i = emit.last; continue; } diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index f5af15de61c..77086b888f6 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -313,6 +313,29 @@ struct ggml_cuda_norm_mxfp4_store { } }; +template +struct ggml_cuda_norm_banked_store { + static constexpr int width = 4; + Store write; + uint3 cols; + uint3 rows; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float4 value) const { + dst[col] = value.x; dst[col + 1] = value.y; dst[col + 2] = value.z; dst[col + 3] = value.w; + const uint2 rc = fast_div_modulo(uint32_t(dst - base + col), cols); + const uint2 bank_row = fast_div_modulo(rc.x, rows); + Store bank = write; + bank.image += int64_t(bank_row.x)*(write.padded/Block)*rows.z; + bank.store(rc.y, bank_row.y, value); + if (rc.y >= cols.z - Group) { + for (int64_t tail = cols.z; tail < write.padded; tail += Group) { + if constexpr (Native) { bank.zero(tail + rc.y%Group, bank_row.y); } + else { bank.store(tail + rc.y%Group, bank_row.y, make_float4(0.0f, 0.0f, 0.0f, 0.0f)); } + } + } + } +}; + struct ggml_cuda_norm_load { __device__ __forceinline__ float operator()(const float * x, int col) const { return x[col]; } }; @@ -689,6 +712,44 @@ static __global__ void rms_norm_mxfp4_f32(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_mxfp4_store{(block_fp4_mmq *) image, cols, padded, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); } +template +static __global__ void rms_norm_banked_mmq_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + void * image, + uint3 cols, + int64_t padded, + uint3 rows, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + if constexpr (Layout == 3) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_banked_store{{(block_fp4_mmq *) image, cols.z, padded, rows.z}, cols, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); + } else { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_banked_store, QK8_1_MMQ, Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32, false>{{(block_q8_1_mmq *) image, cols.z, padded, rows.z}, cols, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); + } +} + template static __global__ void rms_norm_back_f32( const float * grad, const float * xf, float * dst, const int ncols, const float eps) { @@ -1858,7 +1919,7 @@ void ggml_cuda_op_hc_post_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_ten template static void ggml_cuda_rms_norm_mmq_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { const ggml_tensor * x = norm->src[0]; const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; @@ -1868,39 +1929,43 @@ static void ggml_cuda_rms_norm_mmq_launch(ggml_backend_cuda_context & ctx, ggml_ const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; const ggml_cuda_kernel_launch_params params = { dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; - ggml_cuda_kernel_launch(rms_norm_mmq_f32, params, + const auto launch = [&](auto kernel, auto image_cols, auto image_rows) { + ggml_cuda_kernel_launch(kernel, params, (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), - ggml_get_op_params_f32(norm, 0), image, cols, padded, rows, + ggml_get_op_params_f32(norm, 0), image, image_cols, padded, image_rows, pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); + }; + if (banked) { launch(rms_norm_banked_mmq_f32, init_fastdiv_values(cols), init_fastdiv_values(rows)); } + else { launch(rms_norm_mmq_f32, cols, rows); } } template static void ggml_cuda_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { if (norm->src[0]->ne[0] < 1024) { - ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows); + ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } else { - ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows); + ggml_cuda_rms_norm_mmq_launch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } } template static void ggml_cuda_rms_norm_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); - if (add) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else if (mul) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else if (scale) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows); } + if (add) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else if (mul) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else if (scale) { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else { ggml_cuda_rms_norm_mmq(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } } template static void ggml_cuda_rms_norm_mxfp4_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { const ggml_tensor * x = norm->src[0]; const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; @@ -1910,42 +1975,46 @@ static void ggml_cuda_rms_norm_mxfp4_launch(ggml_backend_cuda_context & ctx, ggm const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; const ggml_cuda_kernel_launch_params params = { dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; - ggml_cuda_kernel_launch(rms_norm_mxfp4_f32< Block, Multiply, Add, Scale>, params, + const auto launch = [&](auto kernel, auto image_cols, auto image_rows) { + ggml_cuda_kernel_launch(kernel, params, (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), - ggml_get_op_params_f32(norm, 0), image, cols, padded, rows, + ggml_get_op_params_f32(norm, 0), image, image_cols, padded, image_rows, pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); + }; + if (banked) { launch(rms_norm_banked_mmq_f32<3, Block, Multiply, Add, Scale>, init_fastdiv_values(cols), init_fastdiv_values(rows)); } + else { launch(rms_norm_mxfp4_f32< Block, Multiply, Add, Scale>, cols, rows); } } template static void ggml_cuda_rms_norm_mxfp4(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { if (norm->src[0]->ne[0] < 1024) { - ggml_cuda_rms_norm_mxfp4_launch< 256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows); + ggml_cuda_rms_norm_mxfp4_launch< 256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } else { - ggml_cuda_rms_norm_mxfp4_launch< 1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows); + ggml_cuda_rms_norm_mxfp4_launch< 1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } } static void ggml_cuda_rms_norm_mxfp4_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, bool banked) { GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); - if (add) { ggml_cuda_rms_norm_mxfp4< true, true, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else if (mul) { ggml_cuda_rms_norm_mxfp4< true, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else if (scale) { ggml_cuda_rms_norm_mxfp4< false, false, true>(ctx, norm, mul, add, scale, image, cols, padded, rows); } - else { ggml_cuda_rms_norm_mxfp4< false, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows); } + if (add) { ggml_cuda_rms_norm_mxfp4< true, true, false>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else if (mul) { ggml_cuda_rms_norm_mxfp4< true, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else if (scale) { ggml_cuda_rms_norm_mxfp4< false, false, true>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } + else { ggml_cuda_rms_norm_mxfp4< false, false, false>(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); } } void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, - ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout) { + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout, bool banked) { switch (layout) { - case MMQ_Q8_1_DS_LAYOUT_D4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; - case MMQ_Q8_1_DS_LAYOUT_DS4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; - case MMQ_Q8_1_DS_LAYOUT_D2S6: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; - case 3: ggml_cuda_rms_norm_mxfp4_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_D4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); break; + case MMQ_Q8_1_DS_LAYOUT_DS4: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: ggml_cuda_rms_norm_mmq_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); break; + case 3: ggml_cuda_rms_norm_mxfp4_dispatch(ctx, norm, mul, add, scale, image, cols, padded, rows, banked); break; default: GGML_ABORT("unsupported MMQ image layout"); } } diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 21d3ebb4487..6ae27b7b38e 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -40,7 +40,7 @@ void ggml_cuda_op_rms_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor void ggml_cuda_op_hc_post_norm_emit_q8(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, void * f16, void * bf16, void * image, int64_t cols, int64_t padded); -void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout); +void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * image, int64_t cols, int64_t padded, int64_t rows, int layout, bool banked = false); bool ggml_cuda_should_fuse_hc_post_norm_scale(const ggml_tensor * post, const ggml_tensor * norm, const ggml_tensor * scale); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index bac536ef1d6..78451c7e753 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5192,10 +5192,12 @@ struct test_mul_mat : public test_case { const int64_t pad; // bytes after the m_v rows of each batch of a, so nb[2] of a is not a multiple of nb[1] const bool repeat_rhs; + const bool norm_rhs; std::string vars() override { std::string s = VARS_TO_STR13(type_a, type_b, m, n, k, bs, nr, per, k_v, o, src_overlap, m_v, pad); if (repeat_rhs) { s += "," + VAR_TO_STR(repeat_rhs); } + if (norm_rhs) { s += "," + VAR_TO_STR(norm_rhs); } return s; } @@ -5227,8 +5229,8 @@ struct test_mul_mat : public test_case { std::array bs = {10, 10}, std::array nr = {2, 2}, std::array per = {0, 1, 2, 3}, - int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0, bool repeat_rhs = false) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad), repeat_rhs(repeat_rhs) {} + int64_t k_v = 0, uint32_t o = 1, bool src_overlap = false, int64_t m_v = 0, int64_t pad = 0, bool repeat_rhs = false, bool norm_rhs = false) + : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o), src_overlap(src_overlap), m_v(m_v), pad(pad), repeat_rhs(repeat_rhs), norm_rhs(norm_rhs) {} ggml_tensor * build_graph(ggml_context * ctx) override { // C^T = A * B^T: (k, m) * (k, n) => (m, n) @@ -5310,6 +5312,10 @@ struct test_mul_mat : public test_case { ggml_set_name(b, "b"); } + if (norm_rhs) { + GGML_ASSERT(type_b == GGML_TYPE_F32 && !repeat_rhs); + b = ggml_rms_norm(ctx, b, 1e-6f); + } ggml_tensor * out = ggml_mul_mat(ctx, a, b); ggml_set_name(out, "out"); for (uint32_t i = 1; i < o; ++i) { @@ -5321,7 +5327,7 @@ struct test_mul_mat : public test_case { return out; } - bool run_whole_graph() override { return o > 1 || repeat_rhs; } + bool run_whole_graph() override { return o > 1 || repeat_rhs || norm_rhs; } std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); @@ -10460,6 +10466,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256)); + for (ggml_type type : {GGML_TYPE_Q4_K, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_MXFP4, GGML_TYPE_NVFP4}) { + for (int tokens : {1, 17}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 65, tokens, 512, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 2, false, 0, 0, false, true)); + } + } + for (ggml_type type : {GGML_TYPE_Q4_K, GGML_TYPE_Q4_0, GGML_TYPE_MXFP4, GGML_TYPE_NVFP4}) { for (int tokens : {1, 17}) { test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 65, tokens, 512, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 2)); From 15b12ec519711b001b6485c317e39478b5109be9 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Thu, 8 Oct 2026 02:08:58 -0400 Subject: [PATCH 24/27] cuda : emit banked HC Q8 activation images Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 2 +- tests/test-backend-ops.cpp | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 07b31f07f68..d8d551596da 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4155,7 +4155,7 @@ static std::vector ggml_cuda_plan_norm_q8(ggml_cgraph * const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if ((match.post && (input->ne[2] != 1 || input->ne[3] != 1)) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || + if (prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } if (match.post && input->ne[0] % MATRIX_ROW_PADDING && match.norm->ne[0] >= input->ne[0]) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 78451c7e753..fb614bbb528 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5335,6 +5335,26 @@ struct test_mul_mat : public test_case { } }; +struct test_mul_mat_hc_norm : public test_mul_mat { + test_mul_mat_hc_norm(std::array banks) + : test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 65, 4, 512, banks, {1, 1}, {0, 1, 2, 3}, 0, 2) {} + + std::string op_desc(ggml_tensor *) override { return "MUL_MAT_HC_NORM"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t tokens = bs[0]*bs[1]; + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, tokens); + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, n, tokens); + ggml_tensor * post = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n, tokens); + ggml_tensor * hc = ggml_dsv4_hc_post(ctx, x, residual, post, nullptr); + ggml_set_output(hc); + ggml_tensor * b = ggml_rms_norm(ctx, ggml_reshape_4d(ctx, hc, k, n, bs[0], bs[1]), 1e-6f); + ggml_set_output(b); + ggml_tensor * a = ggml_new_tensor_2d(ctx, type_a, k, m); + return ggml_add(ctx, ggml_mul_mat(ctx, a, b), ggml_mul_mat(ctx, a, b)); + } +}; + // GGML_HINT_SRC0_IS_HADAMARD struct test_mul_mat_hadamard : public test_mul_mat { test_mul_mat_hadamard(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, @@ -10690,6 +10710,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 45, 64, { 8, 1}, {4, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1056, 1, 193, {1, 1}, {4, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1056, 1, 67, {1, 1}, {4, 1}, {0, 2, 1, 3})); + test_cases.emplace_back(new test_mul_mat_hc_norm({3, 1})); + test_cases.emplace_back(new test_mul_mat_hc_norm({1, 3})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 16, 32, 32, { 1, 1}, {1, 1}, {0, 1, 2, 3}, 64, 3)); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 77, 77, {12,1}, {1,1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 4, 96, {3, 2}, {1, 1}, {0, 1, 2, 3}, 0, 1, true)); From 5aab048625d3039ea35e0ba73dc38c8e7cd00a95 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Thu, 8 Oct 2026 03:05:33 -0400 Subject: [PATCH 25/27] cuda : emit banked HC MMQ activation images Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 6 +-- ggml/src/ggml-cuda/norm.cu | 85 ++++++++++++++++++++++++++------- ggml/src/ggml-cuda/norm.cuh | 3 +- tests/test-backend-ops.cpp | 8 +++- 4 files changed, 79 insertions(+), 23 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index d8d551596da..1aea12737b0 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4319,7 +4319,7 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph const ggml_tensor * input = graph->nodes[node]->src[1]; const ggml_tensor * root = input->view_src ? input->view_src : input; uintptr_t begin, end; - if ((match.post && (input->ne[2] != 1 || input->ne[3] != 1)) || keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || + if (keys[node] < 0 || keys[node] > 3 || (match.post && keys[node] > 2) || norm_group_invalid(node) || prepare <= match.last || root != match.dst || !ggml_is_contiguous(input) || input->data != match.dst->data || input->type != GGML_TYPE_F32 || input->ne[0] % QK8_1 || ggml_nbytes(input) != ggml_nbytes(match.dst) || !ggml_cuda_prepared_range(input, device, begin, end)) { return false; } @@ -7037,10 +7037,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud void * f16 = typed ? image(typed->f16) : nullptr; void * bf16 = typed ? image(typed->bf16) : nullptr; if (affine_emit) { - launch_affine_emit({f16, bf16, nullptr, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]}); + launch_affine_emit({f16, bf16, nullptr, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1}); } else { ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, - mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader]); + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); } i = emit.last; continue; diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 77086b888f6..22fa21ac72f 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1477,6 +1477,24 @@ struct ggml_cuda_hc_norm_emit_mmq_pair_store { } }; +template +struct ggml_cuda_hc_norm_banked_pair_store { + static constexpr int width = Store::width; + Store write; + uint3 rows; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float2 value) const { + const uint2 rc = fast_div_modulo(uint32_t(dst - base + col), write.cols); + const uint2 bank_row = fast_div_modulo(rc.x, rows); + const int64_t offset = int64_t(bank_row.x)*write.cols.z*rows.z; + Store bank = write; + bank.image += int64_t(bank_row.x)*(write.padded/QK8_1_MMQ)*rows.z; + if (bank.f16) { bank.f16 += offset; } + if (bank.bf16) { bank.bf16 += offset; } + bank(dst, base + offset, col, value); + } +}; + struct hc_post_norm_data { const float * x; const float * residual; @@ -1619,11 +1637,25 @@ static __global__ __launch_bounds__(Block) void hc_post_norm_emit_mmq_f32(hc_pos ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) image, cols, padded, rows, f16, bf16, scale}); } -template +template +static __global__ __launch_bounds__(Block) void hc_post_norm_emit_banked_mmq_f32(hc_post_norm_data a, + half * f16, nv_bfloat16 * bf16, void * image, uint3 cols, int64_t padded, uint3 rows, float scale) { + hc_post_norm_f32_impl(a, + ggml_cuda_hc_norm_banked_pair_store>{ + {(block_q8_1_mmq *) image, cols, padded, rows.z, f16, bf16, scale}, rows}); +} + +template static auto hc_post_norm_emit_mmq_kernel(int layout) { - if (layout == 1) { return hc_post_norm_emit_mmq_f32; } - if (layout == 2) { return hc_post_norm_emit_mmq_f32; } - return hc_post_norm_emit_mmq_f32; + if constexpr (Banked) { + if (layout == 1) { return hc_post_norm_emit_banked_mmq_f32; } + if (layout == 2) { return hc_post_norm_emit_banked_mmq_f32; } + return hc_post_norm_emit_banked_mmq_f32; + } else { + if (layout == 1) { return hc_post_norm_emit_mmq_f32; } + if (layout == 2) { return hc_post_norm_emit_mmq_f32; } + return hc_post_norm_emit_mmq_f32; + } } struct ggml_cuda_norm_scale_store { @@ -2019,7 +2051,7 @@ void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * no } } -template +template static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows) { const ggml_tensor * x = post->src[0]; const ggml_tensor * residual = post->src[1]; @@ -2056,24 +2088,37 @@ static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; + const auto emit = [&](auto kernel, float factor) { + if constexpr (Banked) { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, init_fastdiv_values(rows), factor); } + else { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, factor); } + }; if (scale) { auto kernel = block_size == 256 - ? (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) - : (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)); - ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, ggml_get_op_params_f32(scale, 0)); + ? (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) + : (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)); + emit(kernel, ggml_get_op_params_f32(scale, 0)); } else { auto kernel = block_size == 256 - ? (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))) - : (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))); - ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, 1.0f); + ? (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))) + : (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))); + emit(kernel, 1.0f); } } -void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout) { +void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout, bool banked) { switch (layout) { - case MMQ_Q8_1_DS_LAYOUT_D4: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; - case MMQ_Q8_1_DS_LAYOUT_DS4: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; - case MMQ_Q8_1_DS_LAYOUT_D2S6: ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); break; + case MMQ_Q8_1_DS_LAYOUT_D4: + if (banked) { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + else { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + if (banked) { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + else { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + if (banked) { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + else { ggml_cuda_hc_post_norm_emit_mmq_dispatch(ctx, post, norm, mul, scale, f16, bf16, image, cols, padded, rows); } + break; default: GGML_ABORT("unsupported MMQ image layout"); } } @@ -2588,7 +2633,10 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens void * bf16 = emit ? emit->bf16 : nullptr; if (emit && emit->mmq) { const uint3 cols = init_fastdiv_values(emit->cols); - const auto launch = [&](auto write) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); }; + const auto launch = [&](auto write) { + if (emit->banked) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, ggml_cuda_hc_norm_banked_pair_store{write, init_fastdiv_values(emit->rows)}); } + else { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); } + }; switch (emit->layout) { case MMQ_Q8_1_DS_LAYOUT_D4: launch(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}); break; case MMQ_Q8_1_DS_LAYOUT_DS4: launch(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}); break; @@ -2607,7 +2655,10 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens GGML_ASSERT(norm); const uint3 cols = init_fastdiv_values(emit->cols); if (emit->mmq) { - const auto launch_mmq = [&](auto write) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); }; + const auto launch_mmq = [&](auto write) { + if (emit->banked) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_hc_norm_banked_pair_store{write, init_fastdiv_values(emit->rows)}); } + else { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); } + }; switch (emit->layout) { case MMQ_Q8_1_DS_LAYOUT_D4: launch_mmq(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}); break; case MMQ_Q8_1_DS_LAYOUT_DS4: launch_mmq(ggml_cuda_hc_norm_emit_mmq_pair_store{(block_q8_1_mmq *) emit->mmq, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}); break; diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 6ae27b7b38e..e66f6cf6e08 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -47,7 +47,7 @@ bool ggml_cuda_should_fuse_hc_post_norm_scale(const ggml_tensor * post, const gg void ggml_cuda_op_hc_post_norm_scale(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * scale, void * f16, void * bf16); -void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout); +void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, int layout, bool banked = false); void ggml_cuda_op_hc_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul); @@ -60,6 +60,7 @@ struct ggml_cuda_hc_affine_emit_data { int64_t padded = 0; int64_t rows = 0; int layout = 0; + bool banked = false; }; void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit = nullptr, ggml_tensor * scale = nullptr); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index fb614bbb528..aba47fd172c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5336,8 +5336,8 @@ struct test_mul_mat : public test_case { }; struct test_mul_mat_hc_norm : public test_mul_mat { - test_mul_mat_hc_norm(std::array banks) - : test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 65, 4, 512, banks, {1, 1}, {0, 1, 2, 3}, 0, 2) {} + test_mul_mat_hc_norm(std::array banks, ggml_type type = GGML_TYPE_Q8_0, int64_t streams = 4) + : test_mul_mat(type, GGML_TYPE_F32, 65, streams, 512, banks, {1, 1}, {0, 1, 2, 3}, 0, 2) {} std::string op_desc(ggml_tensor *) override { return "MUL_MAT_HC_NORM"; } @@ -10712,6 +10712,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1056, 1, 67, {1, 1}, {4, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat_hc_norm({3, 1})); test_cases.emplace_back(new test_mul_mat_hc_norm({1, 3})); + test_cases.emplace_back(new test_mul_mat_hc_norm({3, 1}, GGML_TYPE_Q8_0, 16)); + test_cases.emplace_back(new test_mul_mat_hc_norm({3, 1}, GGML_TYPE_Q4_0, 16)); + test_cases.emplace_back(new test_mul_mat_hc_norm({3, 1}, GGML_TYPE_Q2_K, 16)); + test_cases.emplace_back(new test_mul_mat_hc_norm({1, 3}, GGML_TYPE_Q8_0, 17)); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 16, 32, 32, { 1, 1}, {1, 1}, {0, 1, 2, 3}, 64, 3)); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 77, 77, {12,1}, {1,1})); From fd417c6d6b66b63081f06d4e212f30208d2b8936 Mon Sep 17 00:00:00 2001 From: Rocco A Date: Sun, 11 Oct 2026 19:41:10 -0400 Subject: [PATCH 26/27] cuda : reuse multiple-layout producer prerequisite for banked HC Adapt the published PR133 and PR130 bundle helpers to the PR152 parent. Keep the original banked single-image API and writer unchanged. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 101 ++++++++++++++++- ggml/src/ggml-cuda/norm.cu | 193 +++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/norm.cuh | 5 + 3 files changed, 279 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1aea12737b0..172c37e46d9 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4245,6 +4245,7 @@ struct ggml_cuda_norm_mmq_match { ggml_tensor * dst = nullptr; int last = -1; int image = -1; + std::array images = {-1, -1, -1}; }; static ggml_cuda_norm_mmq_match ggml_cuda_match_norm_mmq(ggml_cgraph * graph, int i, int device) { @@ -4297,7 +4298,7 @@ static ggml_cuda_norm_mmq_match ggml_cuda_match_norm_mmq(ggml_cgraph * graph, in } static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph * graph, int device, - const std::vector & keys, const std::vector & sizes, ggml_cuda_reuse_plan & reuse) { + const std::vector & keys, const std::vector & sizes, ggml_cuda_reuse_plan & reuse, const std::vector & q8_emits, const std::vector & norm_emits) { std::unordered_map> readers; for (int i = 0; i < graph->n_nodes; ++i) { if (keys[i] < 0) { continue; } @@ -4310,6 +4311,7 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph std::vector removed(reuse.groups.size(), false); for (int i = 0; i < graph->n_nodes; ++i) { auto match = ggml_cuda_match_norm_mmq(graph, i, device); + const int prepare = i; if (!match.norm) { continue; } const auto norm_group_invalid = [&](int node) { const int group = keys[node] == int(MMQ_Q8_1_DS_LAYOUT_D2S6) ? 64 : 32; @@ -4376,6 +4378,60 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph moved = true; } } + const auto & info = ggml_cuda_info().devices[device]; + bool small_grid = GGML_CUDA_CC_IS_NVIDIA(info.cc) && info.warp_size == 32 && info.nsm > 0; + uint64_t grid = 1; + if (match.norm) { + for (int d = 1; d < GGML_MAX_DIMS && small_grid; ++d) { + const uint64_t count = match.norm->src[0]->ne[d]; + if (!count || count > uint64_t(info.nsm)/grid) { small_grid = false; break; } + grid *= count; + } + } + const auto * typed = !norm_emits.empty() && norm_emits[i].norm ? &norm_emits[i] : nullptr; + const bool compatible_typed = !typed || (match.post && typed->post == match.post && + typed->norm == match.norm && typed->dst == match.dst && typed->last == match.last); + const int block = match.norm && match.norm->ne[0] < 1024 ? 256 : 1024; + const bool bounded_post = !match.post || match.norm->ne[0] <= 6*block; + if (match.norm && match.image >= 0 && small_grid && bounded_post && compatible_typed && + (q8_emits.empty() || !q8_emits[i].norm)) { + const auto & primary = reuse.groups[match.image]; + const ggml_tensor * shape = graph->nodes[primary.node]->src[1]; + const int primary_key = keys[primary.node]; + if (primary_key >= 0 && primary_key <= 2 && shape->ne[2] == 1 && shape->ne[3] == 1) { + match.images[primary_key] = match.image; + if (found != readers.end()) { + for (const int node : found->second) { + const int key = keys[node]; + if (key < 0 || key > 2 || sizes[node] > SIZE_MAX - 255 || !can_emit(node, node) || + !ggml_are_same_shape(shape, graph->nodes[node]->src[1])) { continue; } + int & target = match.images[key]; + const int previous = reuse.nodes[node]; + if (target >= 0 && previous == target) { continue; } + if (previous >= 0 && !can_emit(reuse.groups[previous].node, reuse.groups[previous].prepare)) { continue; } + if (target < 0) { + target = previous; + if (target < 0) { + target = int(reuse.groups.size()); + reuse.groups.push_back({node, prepare, node, -1, true, GGML_PAD(sizes[node], 256)}); + removed.push_back(false); + } + reuse.groups[target].prepare = prepare; + reuse.groups[target].after = true; + } else if (previous >= 0) { + reuse.groups[target].last = std::max(reuse.groups[target].last, reuse.groups[previous].last); + reuse.groups[target].size = std::max(reuse.groups[target].size, reuse.groups[previous].size); + for (int & member : reuse.nodes) { if (member == previous) { member = target; } } + removed[previous] = true; + } + reuse.nodes[node] = target; + reuse.groups[target].last = std::max(reuse.groups[target].last, node); + reuse.groups[target].size = std::max(reuse.groups[target].size, GGML_PAD(sizes[node], 256)); + moved = true; + } + } + } + } if (match.image >= 0) { emits[i] = match; } } if (moved) { @@ -4388,7 +4444,10 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph } reuse.groups.resize(next); for (int & group : reuse.nodes) { if (group >= 0) { group = remap[group]; } } - for (auto & emit : emits) { if (emit.image >= 0) { emit.image = remap[emit.image]; } } + for (auto & emit : emits) { + if (emit.image >= 0) { emit.image = remap[emit.image]; } + for (int & image : emit.images) { if (image >= 0) { image = remap[image]; } } + } reuse.pack(); } if (reuse.groups.empty()) { return {}; } @@ -6833,7 +6892,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_reuse_plan mmq_reuse(cgraph, stream_ctx, mmq_keys, mmq_sizes, ggml_cuda_mmq_input_overwritten); std::vector mmq_emits; if (!disable_reuse && stream_ctx.concurrent_events.empty()) { - mmq_emits = ggml_cuda_plan_norm_mmq(cgraph, cuda_ctx->device, mmq_keys, mmq_sizes, mmq_reuse); + mmq_emits = ggml_cuda_plan_norm_mmq(cgraph, cuda_ctx->device, mmq_keys, mmq_sizes, mmq_reuse, q8_emits, norm_emits); } ggml_cuda_reuse_inputs mmq_inputs(cuda_ctx->pool(), mmq_reuse); const auto prepare_mmq_group = [&](int group) { @@ -7036,8 +7095,24 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud }; void * f16 = typed ? image(typed->f16) : nullptr; void * bf16 = typed ? image(typed->bf16) : nullptr; + void * images[3] = {}; + int count = 0; + for (int layout = 0; layout < 3; ++layout) { + const int group = emit.images[layout]; + if (group < 0) { continue; } + if (group != g) { mmq_inputs[group].quantized.alloc(mmq_reuse.groups[group].size); } + images[layout] = mmq_inputs[group].quantized.get(); + ++count; + } if (affine_emit) { - launch_affine_emit({f16, bf16, nullptr, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1}); + ggml_cuda_hc_affine_emit_data data = {f16, bf16, nullptr, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1}; + if (count > 1) { + for (int layout = 0; layout < 3; ++layout) { data.mmq_images[layout] = images[layout]; } + } + launch_affine_emit(data); + } else if (count > 1) { + ggml_cuda_op_hc_post_norm_emit_mmq_multi(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, + f16, bf16, images, input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1]); } else { ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); @@ -7129,8 +7204,22 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud const int reader = mmq_reuse.groups[g].node; const ggml_tensor * input = cgraph->nodes[reader]->src[1]; mmq_inputs[g].quantized.alloc(mmq_reuse.groups[g].size); - ggml_cuda_op_rms_norm_mmq(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, - mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); + void * images[3] = {}; + int count = 0; + for (int layout = 0; layout < 3; ++layout) { + const int group = emit.images[layout]; + if (group < 0) { continue; } + if (group != g) { mmq_inputs[group].quantized.alloc(mmq_reuse.groups[group].size); } + images[layout] = mmq_inputs[group].quantized.get(); + ++count; + } + if (count > 1) { + ggml_cuda_op_rms_norm_mmq_multi(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, + images, input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1]); + } else { + ggml_cuda_op_rms_norm_mmq(*cuda_ctx, emit.norm, emit.mul, emit.add, emit.scale, + mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); + } i = emit.last; continue; } diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 22fa21ac72f..a1d93a3890c 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -226,6 +226,23 @@ struct ggml_cuda_norm_mmq_store { } }; +struct ggml_cuda_norm_mmq_bundle_store { + static constexpr int width = 4; + block_q8_1_mmq * images[3]; + int64_t cols; + int64_t padded; + int64_t rows; + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float4 value) const { + dst[col] = value.x; dst[col + 1] = value.y; + dst[col + 2] = value.z; dst[col + 3] = value.w; + const int64_t index = dst - base + col; + if (images[0]) { ggml_cuda_norm_mmq_store{images[0], cols, padded, rows}.emit(index, value); } + if (images[1]) { ggml_cuda_norm_mmq_store{images[1], cols, padded, rows}.emit(index, value); } + if (images[2]) { ggml_cuda_norm_mmq_store{images[2], cols, padded, rows}.emit(index, value); } + } +}; + struct ggml_cuda_norm_mxfp4_store { static constexpr int width = 4; block_fp4_mmq * image; @@ -679,6 +696,36 @@ static __global__ void rms_norm_mmq_f32(const float * x, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, ggml_cuda_norm_mmq_store{(block_q8_1_mmq *) image, cols, padded, rows}, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); } +template +static __global__ void rms_norm_mmq_multi_f32(const float * x, + float * dst, + const int ncols, + const int64_t stride_row, + const int64_t stride_channel, + const int64_t stride_sample, + const float eps, + const ggml_cuda_norm_mmq_bundle_store write, + const float * mul = nullptr, + const int64_t mul_stride_row = 0, + const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, + const uint3 mul_ncols_packed = make_uint3(0, 0, 0), + const uint3 mul_nrows_packed = make_uint3(0, 0, 0), + const uint3 mul_nchannels_packed = make_uint3(0, 0, 0), + const uint3 mul_nsamples_packed = make_uint3(0, 0, 0), + const float * add = nullptr, + const int64_t add_stride_row = 0, + const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, + const uint3 add_ncols_packed = make_uint3(0, 0, 0), + const uint3 add_nrows_packed = make_uint3(0, 0, 0), + const uint3 add_nchannels_packed = make_uint3(0, 0, 0), + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + const float scale_out = 1.0f) { + rms_norm_f32_impl( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, write, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); +} + template static __global__ void rms_norm_mxfp4_f32(const float * x, float * dst, @@ -1495,6 +1542,49 @@ struct ggml_cuda_hc_norm_banked_pair_store { } }; +template +struct ggml_cuda_hc_norm_mmq_bundle_store { + static constexpr int width = 2; + block_q8_1_mmq * images[3]; + uint3 cols; + int64_t padded; + int64_t rows; + half * f16; + nv_bfloat16 * bf16; + float scale; + + template + __device__ __forceinline__ void emit(uint2 rc, float2 value) const { + constexpr int group = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + const ggml_cuda_hc_norm_emit_mmq_pair_store write = { + images[int(Layout)], cols, padded, rows, nullptr, nullptr, 1.0f}; + write.store(rc.y, rc.x, value); + if (rc.y >= cols.z - group) { + for (int64_t tail = cols.z; tail < padded; tail += group) { + write.store(tail + rc.y % group, rc.x, make_float2(0.0f, 0.0f)); + } + } + } + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float2 value) const { + if constexpr (Scale) { value.x = scale * value.x; value.y = scale * value.y; } + dst[col] = value.x; dst[col + 1] = value.y; + const int64_t index = dst - base + col; + const uint2 rc = fast_div_modulo(uint32_t(index), cols); + if (images[0]) { emit(rc, value); } + if (images[1]) { emit(rc, value); } + if (images[2]) { emit(rc, value); } + if (f16) { + f16[index] = ggml_cuda_cast(value.x); + f16[index + 1] = ggml_cuda_cast(value.y); + } + if (bf16) { + bf16[index] = ggml_cuda_cast(value.x); + bf16[index + 1] = ggml_cuda_cast(value.y); + } + } +}; + struct hc_post_norm_data { const float * x; const float * residual; @@ -1645,9 +1735,18 @@ static __global__ __launch_bounds__(Block) void hc_post_norm_emit_banked_mmq_f32 {(block_q8_1_mmq *) image, cols, padded, rows.z, f16, bf16, scale}, rows}); } -template +template +static __global__ __launch_bounds__(Block) void hc_post_norm_emit_mmq_multi_f32(hc_post_norm_data a, ggml_cuda_hc_norm_mmq_bundle_store write) { + hc_post_norm_f32_impl(a, write); +} + +template static auto hc_post_norm_emit_mmq_kernel(int layout) { - if constexpr (Banked) { + if constexpr (Multi) { + if (layout == 1) { return hc_post_norm_emit_mmq_multi_f32; } + if (layout == 2) { return hc_post_norm_emit_mmq_multi_f32; } + return hc_post_norm_emit_mmq_multi_f32; + } else if constexpr (Banked) { if (layout == 1) { return hc_post_norm_emit_banked_mmq_f32; } if (layout == 2) { return hc_post_norm_emit_banked_mmq_f32; } return hc_post_norm_emit_banked_mmq_f32; @@ -2051,8 +2150,50 @@ void ggml_cuda_op_rms_norm_mmq(ggml_backend_cuda_context & ctx, ggml_tensor * no } } -template -static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows) { +template +static void ggml_cuda_rms_norm_mmq_multi_launch(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * const images[3], int64_t cols, int64_t padded, int64_t rows) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * weight = mul ? (mul->src[0] == norm ? mul->src[1] : mul->src[0]) : nullptr; + const ggml_tensor * bias = add ? (add->src[0] == mul ? add->src[1] : add->src[0]) : nullptr; + ggml_tensor * dst = scale ? scale : add ? add : mul ? mul : norm; + const auto pointer = [](const ggml_tensor * t) { return t ? (const float *) t->data : nullptr; }; + const auto stride = [](const ggml_tensor * t, int d) { return t ? int64_t(t->nb[d]/sizeof(float)) : int64_t(0); }; + const auto divisor = [](const ggml_tensor * t, int d) { return t ? init_fastdiv_values(t->ne[d]) : make_uint3(0, 0, 0); }; + const ggml_cuda_kernel_launch_params params = { + dim3(x->ne[1], x->ne[2], x->ne[3]), dim3(Block, 1, 1), 32*sizeof(float), ctx.stream()}; + const ggml_cuda_norm_mmq_bundle_store write = {{(block_q8_1_mmq *) images[0], (block_q8_1_mmq *) images[1], (block_q8_1_mmq *) images[2]}, cols, padded, rows}; + ggml_cuda_kernel_launch(rms_norm_mmq_multi_f32, params, + (const float *) x->data, (float *) dst->data, int(x->ne[0]), stride(x, 1), stride(x, 2), stride(x, 3), + ggml_get_op_params_f32(norm, 0), write, + pointer(weight), stride(weight, 1), stride(weight, 2), stride(weight, 3), + divisor(weight, 0), divisor(weight, 1), divisor(weight, 2), divisor(weight, 3), + pointer(bias), stride(bias, 1), stride(bias, 2), stride(bias, 3), + divisor(bias, 0), divisor(bias, 1), divisor(bias, 2), divisor(bias, 3), + scale ? ggml_get_op_params_f32(scale, 0) : 1.0f); +} + +template +static void ggml_cuda_rms_norm_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * const images[3], int64_t cols, int64_t padded, int64_t rows) { + if (norm->src[0]->ne[0] < 1024) { + ggml_cuda_rms_norm_mmq_multi_launch<256, Multiply, Add, Scale>(ctx, norm, mul, add, scale, images, cols, padded, rows); + } else { + ggml_cuda_rms_norm_mmq_multi_launch<1024, Multiply, Add, Scale>(ctx, norm, mul, add, scale, images, cols, padded, rows); + } +} + +void ggml_cuda_op_rms_norm_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * norm, + ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * const images[3], int64_t cols, int64_t padded, int64_t rows) { + GGML_ASSERT(ggml_get_op_params_f32(norm, 0) >= 0.0f); + if (add) { ggml_cuda_rms_norm_mmq_multi(ctx, norm, mul, add, scale, images, cols, padded, rows); } + else if (mul) { ggml_cuda_rms_norm_mmq_multi(ctx, norm, mul, add, scale, images, cols, padded, rows); } + else if (scale) { ggml_cuda_rms_norm_mmq_multi(ctx, norm, mul, add, scale, images, cols, padded, rows); } + else { ggml_cuda_rms_norm_mmq_multi(ctx, norm, mul, add, scale, images, cols, padded, rows); } +} + +template +static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * image, int64_t cols, int64_t padded, int64_t rows, void * const images[3] = nullptr) { const ggml_tensor * x = post->src[0]; const ggml_tensor * residual = post->src[1]; const ggml_tensor * weights = post->src[2]; @@ -2088,20 +2229,25 @@ static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & const dim3 grid(norm->ne[1], norm->ne[2], norm->ne[3]); const ggml_cuda_kernel_launch_params launch(grid, dim3(block_size), 32*sizeof(float), ctx.stream()); const int layout = a.ncols == a.n_embd ? 1 : a.ncols == a.n_embd*a.hc ? 2 : 0; - const auto emit = [&](auto kernel, float factor) { - if constexpr (Banked) { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, init_fastdiv_values(rows), factor); } + const auto emit = [&](auto kernel, float factor, auto scaled) { + if constexpr (Multi) { + const ggml_cuda_hc_norm_mmq_bundle_store write = { + {(block_q8_1_mmq *) images[0], (block_q8_1_mmq *) images[1], (block_q8_1_mmq *) images[2]}, + init_fastdiv_values(cols), padded, rows, (half *) f16, (nv_bfloat16 *) bf16, factor}; + ggml_cuda_kernel_launch(kernel, launch, a, write); + } else if constexpr (Banked) { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, init_fastdiv_values(rows), factor); } else { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, factor); } }; if (scale) { auto kernel = block_size == 256 - ? (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) - : (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)); - emit(kernel, ggml_get_op_params_f32(scale, 0)); + ? (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) + : (comb ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)); + emit(kernel, ggml_get_op_params_f32(scale, 0), std::true_type{}); } else { auto kernel = block_size == 256 - ? (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))) - : (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))); - emit(kernel, 1.0f); + ? (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))) + : (comb ? (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout)) : (mul ? hc_post_norm_emit_mmq_kernel(layout) : hc_post_norm_emit_mmq_kernel(layout))); + emit(kernel, 1.0f, std::false_type{}); } } @@ -2123,6 +2269,13 @@ void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_te } } +void ggml_cuda_op_hc_post_norm_emit_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * post, + ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, + void * const images[3], int64_t cols, int64_t padded, int64_t rows) { + ggml_cuda_hc_post_norm_emit_mmq_dispatch( + ctx, post, norm, mul, scale, f16, bf16, nullptr, cols, padded, rows, images); +} + struct hc_injection_data { const float * input; float * pre; @@ -2631,7 +2784,13 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens const float factor = ggml_get_op_params_f32(scale, 0); void * f16 = emit ? emit->f16 : nullptr; void * bf16 = emit ? emit->bf16 : nullptr; - if (emit && emit->mmq) { + if (emit && (emit->mmq_images[0] || emit->mmq_images[1] || emit->mmq_images[2])) { + GGML_ASSERT(!emit->banked); + const ggml_cuda_hc_norm_mmq_bundle_store write = { + {(block_q8_1_mmq *) emit->mmq_images[0], (block_q8_1_mmq *) emit->mmq_images[1], (block_q8_1_mmq *) emit->mmq_images[2]}, + init_fastdiv_values(emit->cols), emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}; + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); + } else if (emit && emit->mmq) { const uint3 cols = init_fastdiv_values(emit->cols); const auto launch = [&](auto write) { if (emit->banked) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, ggml_cuda_hc_norm_banked_pair_store{write, init_fastdiv_values(emit->rows)}); } @@ -2654,7 +2813,13 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens if (emit) { GGML_ASSERT(norm); const uint3 cols = init_fastdiv_values(emit->cols); - if (emit->mmq) { + if (emit->mmq_images[0] || emit->mmq_images[1] || emit->mmq_images[2]) { + GGML_ASSERT(!emit->banked); + const ggml_cuda_hc_norm_mmq_bundle_store write = { + {(block_q8_1_mmq *) emit->mmq_images[0], (block_q8_1_mmq *) emit->mmq_images[1], (block_q8_1_mmq *) emit->mmq_images[2]}, + cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}; + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); + } else if (emit->mmq) { const auto launch_mmq = [&](auto write) { if (emit->banked) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_hc_norm_banked_pair_store{write, init_fastdiv_values(emit->rows)}); } else { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); } diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index e66f6cf6e08..8640577b6b0 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -61,6 +61,11 @@ struct ggml_cuda_hc_affine_emit_data { int64_t rows = 0; int layout = 0; bool banked = false; + void * mmq_images[3] = {}; }; void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tensor * first, ggml_tensor * added, ggml_tensor * unary, ggml_tensor * last, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, const ggml_cuda_hc_affine_emit_data * emit = nullptr, ggml_tensor * scale = nullptr); + +void ggml_cuda_op_rms_norm_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * const images[3], int64_t cols, int64_t padded, int64_t rows); + +void ggml_cuda_op_hc_post_norm_emit_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * const images[3], int64_t cols, int64_t padded, int64_t rows); From 48a2df20ce12eadc2a7085249fd3c51ca7e2271f Mon Sep 17 00:00:00 2001 From: Rocco A Date: Sun, 11 Oct 2026 19:41:10 -0400 Subject: [PATCH 27/27] cuda : emit multiple banked HC MMQ activation layouts Reuse the original normalization and quantization primitives to prepare D4, DS4 and D2S6 images for eligible banked consumers in one producer. Keep existing single-image kernels and conservative dispatch contracts. Assisted-by: Codex --- ggml/src/ggml-cuda/ggml-cuda.cu | 4 +- ggml/src/ggml-cuda/norm.cu | 126 +++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/norm.cuh | 2 +- 3 files changed, 120 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 172c37e46d9..79f94d5c515 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4398,7 +4398,7 @@ static std::vector ggml_cuda_plan_norm_mmq(ggml_cgraph const auto & primary = reuse.groups[match.image]; const ggml_tensor * shape = graph->nodes[primary.node]->src[1]; const int primary_key = keys[primary.node]; - if (primary_key >= 0 && primary_key <= 2 && shape->ne[2] == 1 && shape->ne[3] == 1) { + if (primary_key >= 0 && primary_key <= 2 && ((shape->ne[2] == 1 && shape->ne[3] == 1) || (match.post && shape->ne[0] >= 256))) { match.images[primary_key] = match.image; if (found != readers.end()) { for (const int node : found->second) { @@ -7112,7 +7112,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud launch_affine_emit(data); } else if (count > 1) { ggml_cuda_op_hc_post_norm_emit_mmq_multi(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, - f16, bf16, images, input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1]); + f16, bf16, images, input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], input->ne[2] != 1 || input->ne[3] != 1); } else { ggml_cuda_op_hc_post_norm_emit_mmq(*cuda_ctx, emit.post, emit.norm, emit.mul, emit.scale, f16, bf16, mmq_inputs[g].quantized.get(), input->ne[0], GGML_PAD(input->ne[0], MATRIX_ROW_PADDING), input->ne[1], mmq_keys[reader], input->ne[2] != 1 || input->ne[3] != 1); diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index a1d93a3890c..5ff7bcb6c23 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -1585,6 +1585,66 @@ struct ggml_cuda_hc_norm_mmq_bundle_store { } }; +template +struct ggml_cuda_hc_norm_banked_bundle_store { + static constexpr int width = 4; + ggml_cuda_hc_norm_mmq_bundle_store write; + uint3 rows; + + template + __device__ __forceinline__ void emit(block_q8_1_mmq * image, uint2 rc, float4 value) const { + constexpr int group = Layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + const ggml_cuda_norm_mmq_store store = {image, write.cols.z, write.padded, write.rows}; + store.store(rc.y, rc.x, value); + if (rc.y >= write.cols.z - group) { + for (int64_t tail = write.cols.z; tail < write.padded; tail += group) { + const int64_t column = tail + rc.y % group; + const int64_t ib = (column/QK8_1_MMQ)*write.rows + rc.x; + ((char4 *) image[ib].qs)[(column % QK8_1_MMQ)/4] = make_char4(0, 0, 0, 0); + const int iqs = column % QK8_1_MMQ; + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { + if (iqs % 16 == 0 && iqs < 96) { + image[ib].d2s6[2 + iqs/16] = 0.0f; + if (iqs % 64 == 0) { image[ib].d2s6[iqs/64] = 0.0f; } + } + } else if (iqs % 32 == 0) { + if constexpr (Layout == MMQ_Q8_1_DS_LAYOUT_DS4) { image[ib].ds4[iqs/32] = make_half2(0.0f, 0.0f); } + else { image[ib].d4[iqs/32] = 0.0f; } + } + } + } + } + + __device__ __forceinline__ void operator()(float * dst, const float * base, int col, float4 value) const { + if constexpr (Scale) { + value.x = write.scale * value.x; value.y = write.scale * value.y; + value.z = write.scale * value.z; value.w = write.scale * value.w; + } + dst[col] = value.x; dst[col + 1] = value.y; + dst[col + 2] = value.z; dst[col + 3] = value.w; + const int64_t index = dst - base + col; + const uint2 rc = fast_div_modulo(uint32_t(index), write.cols); + const uint2 bank_row = fast_div_modulo(rc.x, rows); + const int64_t offset = int64_t(bank_row.x)*(write.padded/QK8_1_MMQ)*rows.z; + const uint2 local = make_uint2(bank_row.y, rc.y); + if (write.images[0]) { emit(write.images[0] + offset, local, value); } + if (write.images[1]) { emit(write.images[1] + offset, local, value); } + if (write.images[2]) { emit(write.images[2] + offset, local, value); } + if (write.f16) { + write.f16[index] = ggml_cuda_cast(value.x); + write.f16[index + 1] = ggml_cuda_cast(value.y); + write.f16[index + 2] = ggml_cuda_cast(value.z); + write.f16[index + 3] = ggml_cuda_cast(value.w); + } + if (write.bf16) { + write.bf16[index] = ggml_cuda_cast(value.x); + write.bf16[index + 1] = ggml_cuda_cast(value.y); + write.bf16[index + 2] = ggml_cuda_cast(value.z); + write.bf16[index + 3] = ggml_cuda_cast(value.w); + } + } +}; + struct hc_post_norm_data { const float * x; const float * residual; @@ -1687,7 +1747,22 @@ static __device__ __forceinline__ void hc_post_norm_f32_impl(hc_post_norm_data a tmp = block_reduce(tmp, s_sum); const float mean = tmp / a.ncols; const float scale = rsqrtf(mean + a.eps); - if constexpr (Write::width == 2) { + if constexpr (Write::width == 4) { + for (int64_t col = 4*threadIdx.x; col < a.ncols; col += 4*block_size) { + float4 value; + if constexpr (do_multiply) { + value = make_float4( + scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])], + scale * a.post_dst[offset + col + 1] * a.mul[fastmodulo(col + 1, a.mul_ne[0])], + scale * a.post_dst[offset + col + 2] * a.mul[fastmodulo(col + 2, a.mul_ne[0])], + scale * a.post_dst[offset + col + 3] * a.mul[fastmodulo(col + 3, a.mul_ne[0])]); + } else { + value = make_float4(scale * a.post_dst[offset + col], scale * a.post_dst[offset + col + 1], + scale * a.post_dst[offset + col + 2], scale * a.post_dst[offset + col + 3]); + } + write(a.dst + offset, a.dst, int(col), value); + } + } else if constexpr (Write::width == 2) { for (int64_t col = 2*threadIdx.x; col < a.ncols; col += 2*block_size) { float x, y; if constexpr (do_multiply) { @@ -1740,9 +1815,18 @@ static __global__ __launch_bounds__(Block) void hc_post_norm_emit_mmq_multi_f32( hc_post_norm_f32_impl(a, write); } +template +static __global__ __launch_bounds__(Block) void hc_post_norm_emit_banked_mmq_multi_f32(hc_post_norm_data a, ggml_cuda_hc_norm_banked_bundle_store write) { + hc_post_norm_f32_impl(a, write); +} + template static auto hc_post_norm_emit_mmq_kernel(int layout) { - if constexpr (Multi) { + if constexpr (Multi && Banked) { + if (layout == 1) { return hc_post_norm_emit_banked_mmq_multi_f32; } + if (layout == 2) { return hc_post_norm_emit_banked_mmq_multi_f32; } + return hc_post_norm_emit_banked_mmq_multi_f32; + } else if constexpr (Multi) { if (layout == 1) { return hc_post_norm_emit_mmq_multi_f32; } if (layout == 2) { return hc_post_norm_emit_mmq_multi_f32; } return hc_post_norm_emit_mmq_multi_f32; @@ -2234,7 +2318,9 @@ static void ggml_cuda_hc_post_norm_emit_mmq_dispatch(ggml_backend_cuda_context & const ggml_cuda_hc_norm_mmq_bundle_store write = { {(block_q8_1_mmq *) images[0], (block_q8_1_mmq *) images[1], (block_q8_1_mmq *) images[2]}, init_fastdiv_values(cols), padded, rows, (half *) f16, (nv_bfloat16 *) bf16, factor}; - ggml_cuda_kernel_launch(kernel, launch, a, write); + if constexpr (Banked) { + ggml_cuda_kernel_launch(kernel, launch, a, ggml_cuda_hc_norm_banked_bundle_store{write, init_fastdiv_values(rows)}); + } else { ggml_cuda_kernel_launch(kernel, launch, a, write); } } else if constexpr (Banked) { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, init_fastdiv_values(rows), factor); } else { ggml_cuda_kernel_launch(kernel, launch, a, (half *) f16, (nv_bfloat16 *) bf16, image, init_fastdiv_values(cols), padded, rows, factor); } }; @@ -2271,7 +2357,12 @@ void ggml_cuda_op_hc_post_norm_emit_mmq(ggml_backend_cuda_context & ctx, ggml_te void ggml_cuda_op_hc_post_norm_emit_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, - void * const images[3], int64_t cols, int64_t padded, int64_t rows) { + void * const images[3], int64_t cols, int64_t padded, int64_t rows, bool banked) { + if (banked) { + ggml_cuda_hc_post_norm_emit_mmq_dispatch( + ctx, post, norm, mul, scale, f16, bf16, nullptr, cols, padded, rows, images); + return; + } ggml_cuda_hc_post_norm_emit_mmq_dispatch( ctx, post, norm, mul, scale, f16, bf16, nullptr, cols, padded, rows, images); } @@ -2661,7 +2752,22 @@ static __device__ __forceinline__ void hc_injection_norm_emit_impl(hc_post_norm_ tmp = block_reduce(tmp, s_sum); const float mean = tmp / a.ncols; const float scale = rsqrtf(mean + a.eps); - if constexpr (Write::width == 2) { + if constexpr (Write::width == 4) { + for (int64_t col = 4*threadIdx.x; col < a.ncols; col += 4*block_size) { + float4 value; + if constexpr (do_multiply) { + value = make_float4( + scale * a.post_dst[offset + col] * a.mul[fastmodulo(col, a.mul_ne[0])], + scale * a.post_dst[offset + col + 1] * a.mul[fastmodulo(col + 1, a.mul_ne[0])], + scale * a.post_dst[offset + col + 2] * a.mul[fastmodulo(col + 2, a.mul_ne[0])], + scale * a.post_dst[offset + col + 3] * a.mul[fastmodulo(col + 3, a.mul_ne[0])]); + } else { + value = make_float4(scale * a.post_dst[offset + col], scale * a.post_dst[offset + col + 1], + scale * a.post_dst[offset + col + 2], scale * a.post_dst[offset + col + 3]); + } + write(a.dst + offset, a.dst, int(col), value); + } + } else if constexpr (Write::width == 2) { for (int64_t col = 2*threadIdx.x; col < a.ncols; col += 2*block_size) { float x, y; if constexpr (do_multiply) { @@ -2785,11 +2891,12 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens void * f16 = emit ? emit->f16 : nullptr; void * bf16 = emit ? emit->bf16 : nullptr; if (emit && (emit->mmq_images[0] || emit->mmq_images[1] || emit->mmq_images[2])) { - GGML_ASSERT(!emit->banked); const ggml_cuda_hc_norm_mmq_bundle_store write = { {(block_q8_1_mmq *) emit->mmq_images[0], (block_q8_1_mmq *) emit->mmq_images[1], (block_q8_1_mmq *) emit->mmq_images[2]}, init_fastdiv_values(emit->cols), emit->padded, emit->rows, (half *) f16, (nv_bfloat16 *) bf16, factor}; - ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); + if (emit->banked) { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, ggml_cuda_hc_norm_banked_bundle_store{write, init_fastdiv_values(emit->rows)}); + } else { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, false, write); } } else if (emit && emit->mmq) { const uint3 cols = init_fastdiv_values(emit->cols); const auto launch = [&](auto write) { @@ -2814,11 +2921,12 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens GGML_ASSERT(norm); const uint3 cols = init_fastdiv_values(emit->cols); if (emit->mmq_images[0] || emit->mmq_images[1] || emit->mmq_images[2]) { - GGML_ASSERT(!emit->banked); const ggml_cuda_hc_norm_mmq_bundle_store write = { {(block_q8_1_mmq *) emit->mmq_images[0], (block_q8_1_mmq *) emit->mmq_images[1], (block_q8_1_mmq *) emit->mmq_images[2]}, cols, emit->padded, emit->rows, (half *) emit->f16, (nv_bfloat16 *) emit->bf16, 1.0f}; - ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); + if (emit->banked) { + ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_hc_norm_banked_bundle_store{write, init_fastdiv_values(emit->rows)}); + } else { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, write); } } else if (emit->mmq) { const auto launch_mmq = [&](auto write) { if (emit->banked) { ggml_cuda_hc_affine_emit_launch(ctx, a, g, norm, comb, mul, ggml_cuda_hc_norm_banked_pair_store{write, init_fastdiv_values(emit->rows)}); } diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 8640577b6b0..4b4d4255dc4 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -68,4 +68,4 @@ void ggml_cuda_op_hc_affine_injection(ggml_backend_cuda_context & ctx, ggml_tens void ggml_cuda_op_rms_norm_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * add, ggml_tensor * scale, void * const images[3], int64_t cols, int64_t padded, int64_t rows); -void ggml_cuda_op_hc_post_norm_emit_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * const images[3], int64_t cols, int64_t padded, int64_t rows); +void ggml_cuda_op_hc_post_norm_emit_mmq_multi(ggml_backend_cuda_context & ctx, ggml_tensor * post, ggml_tensor * norm, ggml_tensor * mul, ggml_tensor * scale, void * f16, void * bf16, void * const images[3], int64_t cols, int64_t padded, int64_t rows, bool banked = false);