diff --git a/.github/workflows/fusion.yml b/.github/workflows/fusion.yml new file mode 100644 index 0000000000..ad7d5ab60e --- /dev/null +++ b/.github/workflows/fusion.yml @@ -0,0 +1,67 @@ +name: Fusion + +on: + workflow_dispatch: # allows manual triggering + push: + branches: + - master + paths: [ + '.github/workflows/fusion.yml', + 'ggml/**', + 'tests/fusion/**', + 'tests/test-fusion.cpp' + ] + + pull_request: + types: [opened, synchronize, reopened] + paths: [ + '.github/workflows/fusion.yml', + 'ggml/**', + 'tests/fusion/**', + 'tests/test-fusion.cpp' + ] + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }} + cancel-in-progress: true + +env: + GGML_NLOOP: 3 + GGML_N_THREADS: 1 + LLAMA_ARG_LOG_COLORS: 1 + LLAMA_ARG_LOG_PREFIX: 1 + LLAMA_ARG_LOG_TIMESTAMPS: 1 + +jobs: + # TODO: add jobs for other backends as they adopt the fusion debug API + metal: + runs-on: [self-hosted, macOS, ARM64] + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + + - name: Build + id: cmake_build + run: | + cmake -B build \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLAMA_FATAL_WARNINGS=ON \ + -DLLAMA_OPENSSL=OFF \ + -DGGML_SCHED_NO_REALLOC=ON \ + -DGGML_BLAS=OFF \ + -DGGML_METAL=ON + time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu) + time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu) + + - name: Generate models + id: generate_models + run: | + rm -rf build-ci-models && mkdir -p build-ci-models + ./build/bin/test-llama-archs -o build-ci-models + + - name: Test fusion + id: test_fusion + run: | + ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv diff --git a/ci/run.sh b/ci/run.sh index 5463597274..294cbe57bb 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -334,6 +334,35 @@ function gg_sum_test_llama_archs_tensor_split { gg_printf '```\n' } +# test_llama_archs_models + +function gg_run_test_llama_archs_models { + cd ${SRC} + + set -e + + # TODO: fix and re-enable `test-llama-archs` on OpenVINO + # TODO: the `test-llama-archs` currently does not build on Windows, so we check if the binary exists + if [ -z ${GG_BUILD_OPENVINO} ] && [ -f ./build-ci-release/bin/test-llama-archs ]; then + rm -rf build-ci-models && mkdir -p build-ci-models + + # generate the dummy models used by the model-dependent tests + ./build-ci-release/bin/test-llama-archs -o build-ci-models 2>&1 + fi + + set +e +} + +function gg_sum_test_llama_archs_models { + gg_printf '### %s\n\n' "${ci}" + + gg_printf 'Generates the dummy models used by the model-dependent tests\n' + gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)" + gg_printf '```\n' + gg_printf '%s\n' "$(cat $OUT/${ci}.log)" + gg_printf '```\n' +} + # test_scripts function gg_run_test_scripts { @@ -790,6 +819,7 @@ ret=0 test $ret -eq 0 && gg_run ctest_debug test $ret -eq 0 && gg_run ctest_release +test $ret -eq 0 && gg_run test_llama_archs_models test $ret -eq 0 && gg_run test_llama_archs_tensor_split if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index a661e710a2..e7afdb6957 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -10,6 +10,7 @@ ggml_add_backend_library(ggml-metal ggml-metal-device.cpp ggml-metal-common.cpp ggml-metal-context.m + ggml-metal-fusion.cpp ggml-metal-ops.cpp ggml-metal-tuning.cpp ) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 6f1638a114..05755eb3b2 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -1,4 +1,5 @@ #include "ggml-metal-common.h" +#include "ggml-metal-fusion.h" #include "ggml.h" #include "ggml-impl.h" @@ -390,59 +391,31 @@ static std::vector ggml_metal_graph_optimize_reorder(const std::vectorn_nodes; - enum ggml_op ops[MAX_FUSE]; - std::vector nodes; nodes.reserve(gf->n_nodes); // fuse nodes: // we don't want to make reorders that break fusing, so we first pack all fusable tensors // and perform the reorder over the fused nodes. after the reorder is done, we unfuse + // + // the fusable sequences are declared in the fusion table (ggml-metal-fuse.cpp), so the + // packing here is driven by the same patterns that the op encoders will later use for (int i = 0; i < n; i++) { node_info node = { /*.node =*/ gf->nodes[i], /*.fused =*/ {}, }; - // fuse only ops that start with these operations - // can be expanded when needed - if (node.op() == GGML_OP_ADD || - node.op() == GGML_OP_NORM || - node.op() == GGML_OP_RMS_NORM) { - ops[0] = node.op(); + const int f = ggml_metal_fusion_max(gf, i); - int f = i + 1; - while (f < n && f < i + MAX_FUSE) { - // conservatively allow fusing only these ops - // can be expanded when needed - if (gf->nodes[f]->op != GGML_OP_ADD && - gf->nodes[f]->op != GGML_OP_MUL && - gf->nodes[f]->op != GGML_OP_NORM && - gf->nodes[f]->op != GGML_OP_RMS_NORM) { - break; - } - ops[f - i] = gf->nodes[f]->op; - f++; - } + // add the fused tensors into the node info so we can unfuse them later + for (int k = 1; k < f; k++) { + ++i; - f -= i; - for (; f > 1; f--) { - if (ggml_can_fuse(gf, i, ops, f)) { - break; - } - } - - // add the fused tensors into the node info so we can unfuse them later - for (int k = 1; k < f; k++) { - ++i; - - // the .dst() becomes the last fused tensor - node.add_fused(gf->nodes[i]); - } + // the .dst() becomes the last fused tensor + node.add_fused(gf->nodes[i]); } nodes.push_back(std::move(node)); diff --git a/ggml/src/ggml-metal/ggml-metal-context.h b/ggml/src/ggml-metal/ggml-metal-context.h index abf4b06ed2..b538b1ad20 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.h +++ b/ggml/src/ggml-metal/ggml-metal-context.h @@ -33,6 +33,7 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx); void ggml_metal_set_n_cb (ggml_metal_t ctx, int n_cb); void ggml_metal_set_abort_callback (ggml_metal_t ctx, ggml_abort_callback abort_callback, void * user_data); + bool ggml_metal_supports_family (ggml_metal_t ctx, int family); void ggml_metal_capture_next_compute(ggml_metal_t ctx); diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 6cdc4006bc..bf4fe2dcd5 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -6,6 +6,7 @@ #import "ggml-metal-impl.h" #import "ggml-metal-common.h" #import "ggml-metal-ops.h" +#import "ggml-metal-fusion.h" #import @@ -36,15 +37,12 @@ struct ggml_metal { // additional, inference-time compiled pipelines ggml_metal_pipelines_t pipelines_ext; - bool use_fusion; bool use_concurrency; bool use_graph_optimize; int debug_graph; - int debug_fusion; - // how many times a given op was fused - uint64_t fuse_cnt[GGML_OP_COUNT]; + struct ggml_metal_fusion_info * finfo; // capture state int capture_compute; @@ -139,7 +137,6 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); - res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; { @@ -147,20 +144,19 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { res->debug_graph = val ? atoi(val) : 0; } - { - const char * val = getenv("GGML_METAL_FUSION_DEBUG"); - res->debug_fusion = val ? atoi(val) : 0; - } - res->use_graph_optimize = true; if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { res->use_graph_optimize = false; } - memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); + res->finfo = ggml_metal_device_get_fusion_info(dev); + if (ggml_metal_fusion_info_stats(res->finfo)) { + ggml_metal_fusion_info_labels_init(res->finfo); + res->n_cb = 0; + } - GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); + GGML_LOG_INFO("%s: use fusion = %s\n", __func__, ggml_metal_fusion_info_enabled(res->finfo) ? "true" : "false"); GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); @@ -222,15 +218,18 @@ void ggml_metal_free(ggml_metal_t ctx) { ctx->pipelines_ext = nil; } - if (ctx->debug_fusion > 0) { + if (ggml_metal_fusion_info_debug(ctx->finfo) > 0) { GGML_LOG_DEBUG("%s: fusion stats:\n", __func__); - for (int i = 0; i < GGML_OP_COUNT; i++) { - if (ctx->fuse_cnt[i] == 0) { + + const int n_fusions = ggml_metal_fusion_info_n_fusions(ctx->finfo); + for (int i = 0; i < n_fusions; i++) { + const uint64_t count = ggml_metal_fusion_info_count(ctx->finfo, i); + if (count == 0) { continue; } // note: cannot use ggml_log here - GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_op_name((enum ggml_op) i), ctx->fuse_cnt[i]); + GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_metal_fusion_info_label(ctx->finfo, i), count); } } @@ -481,10 +480,17 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * @autoreleasepool { ctx->gf = gf; - ctx->n_nodes_0 = MIN(n_main, gf->n_nodes); - ctx->n_nodes_1 = gf->n_nodes - ctx->n_nodes_0; + if (ctx->n_cb == 0) { + // single-threaded encoding: the whole graph is encoded by one command buffer + ctx->n_nodes_0 = gf->n_nodes; + ctx->n_nodes_1 = 0; + ctx->n_nodes_per_cb = 0; + } else { + ctx->n_nodes_0 = MIN(n_main, gf->n_nodes); + ctx->n_nodes_1 = gf->n_nodes - ctx->n_nodes_0; - ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb; + ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb; + } if (ctx->capture_compute >= 0) { ctx->capture_compute--; @@ -682,6 +688,12 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) { } void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { + // when fusion stats are collected the graph must be encoded by a single thread so the + // counters are race-free; override whatever the caller requested + if (ggml_metal_fusion_info_stats(ctx->finfo)) { + n_cb = 0; + } + if (ctx->n_cb != n_cb) { ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS); @@ -717,13 +729,12 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { ctx->dev, cmd_buf, ctx->gf, + ctx->finfo, idx_start, idx_end, - ctx->use_fusion, ctx->use_concurrency, ctx->capture_compute, - ctx->debug_graph, - ctx->debug_fusion); + ctx->debug_graph); for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) { const int res = ggml_metal_op_encode(ctx_op, idx); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 31fc07d44d..ced33aadfb 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -325,6 +325,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev); +struct ggml_metal_fusion_info; + +// the device-owned fusion debugging context (NULL unless fusion debugging is enabled) +struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev); + // // device buffers // diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index afd6f52101..5654c50040 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1,4 +1,5 @@ #import "ggml-metal-device.h" +#import "ggml-metal-fusion.h" #import "ggml-impl.h" #import "ggml-backend-impl.h" @@ -896,6 +897,9 @@ struct ggml_metal_device { struct ggml_metal_device_props props; + // shared fusion debugging context + struct ggml_metal_fusion_info * finfo; + // virtual address for GPU memory allocations atomic_uintptr_t addr_virt; }; @@ -1274,6 +1278,13 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) { dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; } + { + const char * val = getenv("GGML_METAL_FUSION_DEBUG"); + dev->finfo = ggml_metal_fusion_info_init( + getenv("GGML_METAL_FUSION_DISABLE") == nil, + val ? atoi(val) : 0); + } + snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); const char * gpu_name = [[dev->mtl_device name] UTF8String]; if (n_devices > 1) { @@ -1348,6 +1359,8 @@ void ggml_metal_device_free(ggml_metal_device_t dev) { assert(dev != NULL); @autoreleasepool { + ggml_metal_fusion_info_free(dev->finfo); + ggml_metal_rsets_free(dev->rsets); ggml_metal_library_free(dev->library); @@ -1935,6 +1948,10 @@ static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev) { dev->props.has_tensor = false; } +struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev) { + return dev->finfo; +} + // // device buffers // diff --git a/ggml/src/ggml-metal/ggml-metal-fusion.cpp b/ggml/src/ggml-metal/ggml-metal-fusion.cpp new file mode 100644 index 0000000000..ac3ac04148 --- /dev/null +++ b/ggml/src/ggml-metal/ggml-metal-fusion.cpp @@ -0,0 +1,502 @@ +#include "ggml-metal-fusion.h" + +#include "ggml-backend-impl.h" +#include "ggml-metal-device.h" + +#include +#include +#include + +// ---- helpers ------------------------------------------------------------- + +// true if two tensors live in the same Metal buffer +static bool ggml_metal_fusion_same_buffer(const ggml_tensor * a, const ggml_tensor * b) { + if (!a || !b) { + return false; + } + + ggml_backend_buffer_t ba = a->view_src ? a->view_src->buffer : a->buffer; + ggml_backend_buffer_t bb = b->view_src ? b->view_src->buffer : b->buffer; + + ggml_metal_buffer_t ca = (ggml_metal_buffer_t) ba->context; + ggml_metal_buffer_t cb = (ggml_metal_buffer_t) bb->context; + + return ggml_metal_buffer_get_id(ca, a).metal == ggml_metal_buffer_get_id(cb, b).metal; +} + +// ---- pattern checks ------------------------------------------------------ + +// NORM/RMS_NORM + MUL + ADD: the weight/bias of each fused step must match the norm input +// width, be contiguous rows, and the fused outputs must stay F32 +static bool ggml_metal_fusion_check_norm( + const ggml_metal_fusion * fusion, + const ggml_tensor * const * nodes, + ggml_metal_fusion_mode mode) { + GGML_UNUSED(mode); + + GGML_ASSERT(fusion->n_ops >= 2); + + for (int j = 1; j < fusion->n_ops; j++) { + // the fused MUL/ADD must read the previous node as src0 + if (nodes[j]->src[0] != nodes[j - 1]) { + return false; + } + + // the weight/bias must have the same row width as the norm input + if (nodes[j]->src[1]->ne[0] != nodes[0]->ne[0]) { + return false; + } + + if (!ggml_is_contiguous_rows(nodes[j]->src[1])) { + return false; + } + + if (nodes[j]->type != GGML_TYPE_F32) { + return false; + } + } + + return true; +} + +// ADD x N: each ADD reads the previous ADD as src0, and all addends must share layout +// (and, in FULL mode, live in the same Metal buffer) +static bool ggml_metal_fusion_check_add_chain( + const ggml_metal_fusion * fusion, + const ggml_tensor * const * nodes, + ggml_metal_fusion_mode mode) { + GGML_ASSERT(fusion->n_ops >= 2); + + for (int j = 1; j < fusion->n_ops; j++) { + if (nodes[j]->src[0] != nodes[j - 1]) { + return false; + } + + if (!ggml_are_same_layout(nodes[j]->src[1], nodes[j - 1]->src[1])) { + return false; + } + + if (mode == GGML_METAL_FUSION_FULL) { + if (!ggml_metal_fusion_same_buffer(nodes[j]->src[1], nodes[0]->src[1])) { + return false; + } + } + } + + return true; +} + +// GATED_DELTA_NET + CPY: the trailing cpy scatters the gdn state snapshots into the recurrent +// cache, so the gdn kernel writes them straight to the cache and the cpy is elided. +// mirrors ggml_metal_op_can_fuse_gdn_cache (PR #25788). the gdn output has other consumers (the +// attn scores view), so unlike the other patterns this is not an elision chain: the structural +// checks live entirely in this callback (unsafe = true). +static bool ggml_metal_fusion_check_gdn_cache( + const ggml_metal_fusion * fusion, + const ggml_tensor * const * nodes, + ggml_metal_fusion_mode mode) { + GGML_UNUSED(fusion); + + const ggml_tensor * gdn = nodes[0]; + const ggml_tensor * cpy = nodes[1]; + + // the kernel skips the snapshot tail, so the gdn output must not be a graph output + if (gdn->type != GGML_TYPE_F32 || (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + + if (cpy->op != GGML_OP_CPY || (cpy->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return false; + } + + const int64_t S_v = gdn->src[2]->ne[0]; + const int64_t H = gdn->src[2]->ne[1]; + const int64_t n_tokens = gdn->src[2]->ne[2]; + const int64_t n_seqs = gdn->src[2]->ne[3]; + const int64_t K = ggml_get_op_params_i32(gdn, 0); + const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs); + + const int64_t D = S_v * S_v * H; + const int64_t n_written = std::min(n_tokens, K); + + const ggml_tensor * src = cpy->src[0]; // gdn snapshot tail view + const ggml_tensor * dst = cpy->src[1]; // cache view + + // src must be this gdn's snapshot tail (contiguous, at the tail offset) + if (src->op != GGML_OP_VIEW || src->view_src != gdn || + src->view_offs != tail_off || !ggml_is_contiguous(src)) { + return false; + } + + const int64_t expected_ne[GGML_MAX_DIMS] = { D, n_seqs, n_written, 1 }; + if (dst->type != GGML_TYPE_F32 || + !std::equal(expected_ne, expected_ne + GGML_MAX_DIMS, dst->ne) || + dst->nb[0] != ggml_type_size(GGML_TYPE_F32) || + dst->nb[1] != ggml_row_size(GGML_TYPE_F32, D)) { + return false; + } + + if (mode == GGML_METAL_FUSION_FULL) { + // the cache must be allocated so the kernel can write straight to its buffer + if (dst->data == nullptr) { + return false; + } + } + + return true; +} + +// MUL + SIN + SQR + MUL + ADD (snake activation) +static bool ggml_metal_fusion_check_snake( + const ggml_metal_fusion * fusion, + const ggml_tensor * const * nodes, + ggml_metal_fusion_mode mode) { + GGML_UNUSED(fusion); + GGML_UNUSED(mode); + + const ggml_tensor * mul0 = nodes[0]; + const ggml_tensor * sin_node = nodes[1]; + const ggml_tensor * sqr = nodes[2]; + const ggml_tensor * mul1 = nodes[3]; + const ggml_tensor * add = nodes[4]; + + // x carries the full activation shape, a is the broadcast operand + const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; + const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; + + // mul1 reads sqr and inv_b in either operand order + const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; + + // closure check: the trailing add reads the same x as the leading mul + const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; + + // x is in the supported whitelist and every chain intermediate shares x's type. + // a and inv_b bind as device const float * in the kernel, so they stay F32. + const bool types_ok = + (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && + (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && + (mul0->type == x->type) && (sin_node->type == x->type) && + (sqr->type == x->type) && (mul1->type == x->type) && + (add->type == x->type); + + // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D + const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; + const bool dim_ok = + (x->ne[2] == 1) && (x->ne[3] == 1) && + (add->ne[2] == 1) && (add->ne[3] == 1) && + (a->ne[2] == 1) && (a->ne[3] == 1) && + (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); + + // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous + const bool contig_ok = + ggml_is_contiguous(x) && ggml_is_contiguous(add) && + ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); + + return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; +} + +// ---- patterns ------------------------------------------------------------ + +static const ggml_op ops_norm_mul[] = { GGML_OP_NORM, GGML_OP_MUL }; +static const ggml_op ops_norm_mul_add[] = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }; +static const ggml_op ops_rms_norm_mul[] = { GGML_OP_RMS_NORM, GGML_OP_MUL }; +static const ggml_op ops_rms_norm_mul_add[] = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }; + +static const ggml_op ops_add_2[] = { GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_add_3[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_add_4[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_add_5[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_add_6[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_add_7[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; + +static const ggml_op ops_gdn_cache[] = { GGML_OP_GATED_DELTA_NET, GGML_OP_CPY }; + +static const ggml_metal_fusion ggml_metal_fusions[] = { + { GGML_METAL_FUSION_NORM_MUL, ops_norm_mul, 2, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add, 3, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL, ops_rms_norm_mul, 2, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, false, ggml_metal_fusion_check_norm }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_2, 2, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_3, 3, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_4, 4, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_5, 5, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_6, 6, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_ADD_CHAIN, ops_add_7, 7, false, ggml_metal_fusion_check_add_chain }, + { GGML_METAL_FUSION_SNAKE, ops_snake, 5, false, ggml_metal_fusion_check_snake }, + { GGML_METAL_FUSION_GDN_CACHE, ops_gdn_cache, 2, true, ggml_metal_fusion_check_gdn_cache }, +}; + +const ggml_metal_fusion * ggml_metal_fusion_all(int * n) { + *n = (int) sizeof(ggml_metal_fusions) / sizeof(ggml_metal_fusions[0]); + + return ggml_metal_fusions; +} + +// ---- shared fusion info --------------------------------------------------- + +static std::string ggml_metal_fusion_label(const ggml_metal_fusion * fusion) { + GGML_ASSERT(fusion != nullptr); + + std::string label; + for (int j = 0; j < fusion->n_ops; j++) { + if (j > 0) { + label += '+'; + } + label += ggml_op_name(fusion->ops[j]); + } + return label; +} + +struct ggml_metal_fusion_info { + std::vector labels; + std::vector counts; + bool enabled; + bool stats; + bool labels_set; + int debug; +}; + +struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug) { + ggml_metal_fusion_info * finfo = new ggml_metal_fusion_info; + finfo->enabled = enabled; + finfo->stats = debug > 0; + finfo->labels_set = false; + finfo->debug = debug; + + if (finfo->stats) { + ggml_metal_fusion_info_labels_init(finfo); + } + + return finfo; +} + +void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo) { + delete finfo; +} + +bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo) { + return finfo->enabled; +} + +bool ggml_metal_fusion_info_stats(const struct ggml_metal_fusion_info * finfo) { + return finfo->stats; +} + +int ggml_metal_fusion_info_debug(const struct ggml_metal_fusion_info * finfo) { + return finfo->debug; +} + +int ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo) { + return (int) finfo->labels.size(); +} + +const char * ggml_metal_fusion_info_label(const struct ggml_metal_fusion_info * finfo, int idx) { + GGML_ASSERT(idx >= 0 && idx < (int) finfo->labels.size()); + return finfo->labels[idx].c_str(); +} + +uint64_t ggml_metal_fusion_info_count(const struct ggml_metal_fusion_info * finfo, int idx) { + GGML_ASSERT(idx >= 0 && idx < (int) finfo->counts.size()); + return finfo->counts[idx]; +} + +void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion) { + if (!finfo->stats || fusion == nullptr) { + return; + } + + int n = 0; + const ggml_metal_fusion * all = ggml_metal_fusion_all(&n); + + int idx = -1; + for (int i = 0; i < n; i++) { + if (&all[i] == fusion) { + idx = i; + break; + } + } + + if (idx >= 0 && idx < (int) finfo->counts.size()) { + finfo->counts[idx]++; + } +} + +void ggml_metal_fusion_info_set_enabled(struct ggml_metal_fusion_info * finfo, bool enabled) { + finfo->enabled = enabled; +} + +void ggml_metal_fusion_info_labels_init(struct ggml_metal_fusion_info * finfo) { + if (finfo->labels_set) { + return; + } + + int n = 0; + const ggml_metal_fusion * all = ggml_metal_fusion_all(&n); + + finfo->labels.clear(); + finfo->counts.assign(n, 0); + finfo->labels.reserve(n); + + for (int i = 0; i < n; i++) { + finfo->labels.emplace_back(ggml_metal_fusion_label(&all[i])); + } + + finfo->labels_set = true; +} + +void ggml_metal_fusion_info_stats_init(struct ggml_metal_fusion_info * finfo) { + finfo->stats = true; + ggml_metal_fusion_info_labels_init(finfo); +} + +void ggml_metal_fusion_info_stats_reset(struct ggml_metal_fusion_info * finfo) { + std::fill(finfo->counts.begin(), finfo->counts.end(), 0); +} + +int ggml_metal_fusion_info_stats_get(const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n) { + const int n_fusions = (int) finfo->labels.size(); + + if (labels == nullptr) { + return n_fusions; + } + + const int n_fill = std::min(n, n_fusions); + for (int i = 0; i < n_fill; i++) { + labels[i] = finfo->labels[i].c_str(); + if (counts != nullptr) { + counts[i] = finfo->counts[i]; + } + } + + return n_fill; +} + +// ---- queries ------------------------------------------------------------- + +// find the longest pattern matching the node sequence starting at idx +// (idx is a position in node_idxs, which maps to graph node indices) +const ggml_metal_fusion * ggml_metal_fusion_next( + const ggml_cgraph * gf, + const int * node_idxs, + int n_idxs, + int idx, + ggml_metal_fusion_mode mode, + int * n_out) { + int n = 0; + const ggml_metal_fusion * all = ggml_metal_fusion_all(&n); + + const ggml_metal_fusion * res = nullptr; + int best = 1; + + for (int i = 0; i < n; i++) { + const ggml_metal_fusion * fusion = &all[i]; + + // only look for a longer match than the current best + if (fusion->n_ops <= best) { + continue; + } + if (idx + fusion->n_ops > n_idxs) { + continue; + } + + const ggml_tensor * nodes[GGML_METAL_FUSION_MAX]; + + // the op sequence must match exactly + bool ok = true; + for (int j = 0; j < fusion->n_ops; j++) { + nodes[j] = gf->nodes[node_idxs[idx + j]]; + if (nodes[j]->op != fusion->ops[j]) { + ok = false; + break; + } + } + if (!ok) { + continue; + } + + if (!fusion->unsafe) { + // common element-wise chain constraints: each node reads the previous one, + // and all nodes have the same shape + for (int j = 1; j < fusion->n_ops && ok; j++) { + if (nodes[j]->src[0] != nodes[j - 1] && nodes[j]->src[1] != nodes[j - 1]) { + ok = false; + break; + } + if (!ggml_are_same_shape(nodes[j], nodes[j - 1])) { + ok = false; + break; + } + } + if (!ok) { + continue; + } + + // all current fusions are single-output elision chains, so the last node is the only output + // TODO: multi-output fusions: store pattern-relative offsets in the table and translate them here + int outputs_buf[1]; + outputs_buf[0] = node_idxs[idx + fusion->n_ops - 1]; + + // structural subgraph checks (op sequence, elidable uses, view containment) + if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fusion->n_ops, fusion->ops, outputs_buf, 1)) { + continue; + } + } + + // pattern-specific checks (the sole validator for unsafe patterns) + if (fusion->check && !fusion->check(fusion, nodes, mode)) { + continue; + } + + best = fusion->n_ops; + res = fusion; + } + + *n_out = best; + + return res; +} + +// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that +// could be fused, chaining patterns back-to-back. matching runs on the same filtered (view +// transparent) node sequence that the compute phase uses, so the returned count is the raw index +// span from idx to the last matched node (intermediate views are packed along). +int ggml_metal_fusion_max(const ggml_cgraph * gf, int idx) { + // an empty/view node cannot start a pattern - pack it alone + if (ggml_op_is_empty(gf->nodes[idx]->op) || ggml_is_empty(gf->nodes[idx])) { + return 1; + } + + // collect the non-empty node indices starting at idx + int idxs[GGML_METAL_FUSION_MAX]; + int n_idxs = 0; + for (int i = idx; i < gf->n_nodes && n_idxs < GGML_METAL_FUSION_MAX; i++) { + if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) { + idxs[n_idxs++] = i; + } + } + if (n_idxs == 0) { + return 1; + } + + int total = 0; + int i_f = 0; + + while (i_f < n_idxs && total < GGML_METAL_FUSION_MAX) { + int len = 1; + const ggml_metal_fusion * fusion = ggml_metal_fusion_next(gf, idxs, n_idxs, i_f, GGML_METAL_FUSION_STRUCTURAL, &len); + if (!fusion || total + len > GGML_METAL_FUSION_MAX) { + break; + } + + total += len; + i_f += len; + } + + if (i_f == 0) { + return 1; + } + + // map the matched non-empty nodes back to the raw index span (views are included) + return std::min(GGML_METAL_FUSION_MAX, idxs[i_f - 1] - idx + 1); +} diff --git a/ggml/src/ggml-metal/ggml-metal-fusion.h b/ggml/src/ggml-metal/ggml-metal-fusion.h new file mode 100644 index 0000000000..e8515bdeca --- /dev/null +++ b/ggml/src/ggml-metal/ggml-metal-fusion.h @@ -0,0 +1,104 @@ +// single source of truth for the fusions supported by the Metal backend +// +// every fusable subgraph is declared exactly once as a ggml_metal_fusion entry in +// the table in ggml-metal-fusion.cpp. both the graph optimizer (ggml_metal_fusion_max) +// and the op encoders (ggml_metal_fusion_next) consult this same table, so the two +// phases can never disagree about what can be fused. + +#pragma once + +#include "ggml-impl.h" + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// the maximum number of nodes that can be fused in a single kernel +// (also the maximum length of a packed fusion group during graph optimization) +#define GGML_METAL_FUSION_MAX 16 + +typedef enum ggml_metal_fusion_mode { + // structural checks only; used by the graph optimizer, at which point the graph + // tensors are not allocated yet, so buffer placement cannot be verified + GGML_METAL_FUSION_STRUCTURAL = 0, + // full checks, including buffer placement; used by the op encoders + GGML_METAL_FUSION_FULL, +} ggml_metal_fusion_mode; + +// identifier of each fusion pattern so the op encoders know which kernel to use +typedef enum ggml_metal_fusion_id { + GGML_METAL_FUSION_NONE = 0, + GGML_METAL_FUSION_NORM_MUL, // NORM/RMS_NORM + MUL + GGML_METAL_FUSION_NORM_MUL_ADD, // NORM/RMS_NORM + MUL + ADD + GGML_METAL_FUSION_ADD_CHAIN, // ADD x N (N in [2, 7]) + GGML_METAL_FUSION_SNAKE, // MUL + SIN + SQR + MUL + ADD + GGML_METAL_FUSION_GDN_CACHE, // GATED_DELTA_NET + CPY (write snapshots into the recurrent cache) +} ggml_metal_fusion_id; + +struct ggml_metal_fusion { + ggml_metal_fusion_id id; + + const enum ggml_op * ops; // op sequence (fixed length) + int n_ops; // number of ops + + // if unsafe: the generic chain/shape + ggml_can_fuse_subgraph checks are skipped and the + // check callback below is the sole validator (used for patterns that are not elision chains, + // e.g. the gdn + cache-cpy write-through fusion) + bool unsafe; + + // extra backend constraints on top of ggml_can_fuse_subgraph + // nodes[j] is the j-th node of the pattern + bool (*check)(const struct ggml_metal_fusion * fusion, + const struct ggml_tensor * const * nodes, + ggml_metal_fusion_mode mode); +}; + +typedef struct ggml_metal_fusion ggml_metal_fusion; + +// the single table of all fusions supported by the Metal backend +const ggml_metal_fusion * ggml_metal_fusion_all(int * n); + +// ---- shared fusion info --------------------------------------------------- + +// shared fusion debugging context, owned by the device; newly created backend contexts for that +// device register with it so the fusion counters are race-free and accumulate across contexts. +struct ggml_metal_fusion_info; // defined in ggml-metal-fusion.cpp + +struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug); +void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo); + +bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo); +bool ggml_metal_fusion_info_stats (const struct ggml_metal_fusion_info * finfo); +int ggml_metal_fusion_info_debug (const struct ggml_metal_fusion_info * finfo); + +int ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo); +const char * ggml_metal_fusion_info_label (const struct ggml_metal_fusion_info * finfo, int idx); +uint64_t ggml_metal_fusion_info_count (const struct ggml_metal_fusion_info * finfo, int idx); + +void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion); +void ggml_metal_fusion_info_set_enabled (struct ggml_metal_fusion_info * finfo, bool enabled); + +void ggml_metal_fusion_info_stats_init ( struct ggml_metal_fusion_info * finfo); +void ggml_metal_fusion_info_stats_reset( struct ggml_metal_fusion_info * finfo); +int ggml_metal_fusion_info_stats_get (const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n); +void ggml_metal_fusion_info_labels_init( struct ggml_metal_fusion_info * finfo); + +// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode`. +// returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed. +const ggml_metal_fusion * ggml_metal_fusion_next( + const struct ggml_cgraph * gf, + const int * node_idxs, + int n_idxs, + int idx, + ggml_metal_fusion_mode mode, + int * n_out); + +// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that +// could be fused, chaining patterns back-to-back. returns at least 1. +int ggml_metal_fusion_max(const struct ggml_cgraph * gf, int idx); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 28a9ba101c..7ad21341e4 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -985,6 +985,7 @@ typedef struct { uint64_t nb1; uint64_t nb2; uint64_t nb3; + uint64_t nb_out; // 0 => snapshots are appended after the attn scores (unfused) } ggml_metal_kargs_gated_delta_net; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 3db8bca437..b4e87cb2cc 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -7,6 +7,7 @@ #include "ggml-metal-impl.h" #include "ggml-metal-common.h" #include "ggml-metal-device.h" +#include "ggml-metal-fusion.h" #include "ggml-metal-tuning.h" #include @@ -31,24 +32,22 @@ struct ggml_metal_op { ggml_metal_device_t dev, ggml_metal_cmd_buf_t cmd_buf, ggml_cgraph * gf, + ggml_metal_fusion_info * finfo, int idx_start, int idx_end, - bool use_fusion, bool use_concurrency, bool use_capture, - int debug_graph, - int debug_fusion) { + int debug_graph) { this->dev = dev; this->lib = ggml_metal_device_get_library(dev); this->enc = ggml_metal_encoder_init(cmd_buf, use_concurrency); this->mem_ranges = ggml_mem_ranges_init(debug_graph); + this->finfo = finfo; this->idx_start = idx_start; this->idx_end = idx_end; - this->use_fusion = use_fusion; this->use_concurrency = use_concurrency; this->use_capture = use_capture; this->debug_graph = debug_graph; - this->debug_fusion = debug_fusion; this->gf = gf; idxs.reserve(gf->n_nodes); @@ -78,15 +77,24 @@ struct ggml_metal_op { return ggml_graph_node(gf, idxs[i]); } - bool can_fuse(int i0, const ggml_op * ops, int n_ops) const { - assert(use_fusion); + // consult the fusion table for the longest pattern starting at i0 + // returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes + const ggml_metal_fusion * can_fuse(int i0, enum ggml_metal_fusion_mode mode, int * n_out) const { + assert(use_fusion()); assert(i0 >= 0 && i0 < n_nodes()); - if (i0 + n_ops > n_nodes()) { - return false; - } + return ggml_metal_fusion_next(gf, idxs.data(), (int) idxs.size(), i0, mode, n_out); + } - return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops); + // whether to attempt fusion; the toggle lives in the shared fusion debugging context owned + // by the device (initialized from GGML_METAL_FUSION_DISABLE, overridable by the test) + bool use_fusion() const { + return ggml_metal_fusion_info_enabled(finfo); + } + + // record that a fusion fired, indexed by the matching table entry + void count_fusions(const ggml_metal_fusion * fusion) const { + ggml_metal_fusion_info_count_fusion(finfo, fusion); } ggml_metal_device_t dev; @@ -94,12 +102,13 @@ struct ggml_metal_op { ggml_metal_encoder_t enc; ggml_mem_ranges_t mem_ranges; - bool use_fusion; + // shared fusion debugging context + ggml_metal_fusion_info * finfo; + bool use_concurrency; bool use_capture; int debug_graph; - int debug_fusion; private: ggml_cgraph * gf; @@ -115,24 +124,22 @@ ggml_metal_op_t ggml_metal_op_init( ggml_metal_device_t dev, ggml_metal_cmd_buf_t cmd_buf, ggml_cgraph * gf, + ggml_metal_fusion_info * finfo, int idx_start, int idx_end, - bool use_fusion, bool use_concurrency, bool use_capture, - int debug_graph, - int debug_fusion) { + int debug_graph) { ggml_metal_op_t res = new ggml_metal_op( dev, cmd_buf, gf, + finfo, idx_start, idx_end, - use_fusion, use_concurrency, use_capture, - debug_graph, - debug_fusion); + debug_graph); return res; } @@ -1868,6 +1875,8 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; + const bool use_fusion = ctx->use_fusion(); + const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); @@ -1880,6 +1889,31 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op); + // when fused with the trailing cache cpy, the snapshots are written straight into the + // recurrent cache and the cpy is skipped (see GGML_METAL_FUSION_GDN_CACHE) + ggml_metal_buffer_id bid_out = ggml_metal_get_buffer_id(op); + uint64_t nb_out = 0; + int n_fuse = 1; + + if (use_fusion) { + int n = 1; + const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n); + + if (fusion && fusion->id == GGML_METAL_FUSION_GDN_CACHE) { + const ggml_tensor * dst_cache = ctx->node(idx + 1)->src[1]; // cache view + + bid_out = ggml_metal_get_buffer_id(dst_cache); + nb_out = dst_cache->nb[2]/sizeof(float); + n_fuse = 2; + + ctx->count_fusions(fusion); + + if (debug_fusion > 1) { + GGML_LOG_DEBUG("%s: fuse: GATED_DELTA_NET + CPY\n", __func__); + } + } + } + int ida = 0; ggml_metal_kargs_gated_delta_net args = { @@ -1918,23 +1952,25 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, + /*.nb_out =*/ nb_out, }; ggml_metal_encoder_set_pipeline(enc, pipeline); - ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); // args ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst (attn) + ggml_metal_encoder_set_buffer (enc, bid_out, ida++); // state_out const int nsg = pipeline.nsg; ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1); - return 1; + return n_fuse; } int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { @@ -3718,56 +3754,20 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { return 1; } -// Snake activation autofuse: mul -> sin -> sqr -> mul -> add -static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) { - static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; - - if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) { - return false; - } - - const ggml_tensor * mul0 = ctx->node(idx + 0); - const ggml_tensor * sin_node = ctx->node(idx + 1); - const ggml_tensor * sqr = ctx->node(idx + 2); - const ggml_tensor * mul1 = ctx->node(idx + 3); - const ggml_tensor * add = ctx->node(idx + 4); - - // x carries the full activation shape, a is the broadcast operand - const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; - const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; - - // mul1 reads sqr and inv_b in either operand order - const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; - - // closure check: the trailing add reads the same x as the leading mul - const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; - - // x is in the supported whitelist and every chain intermediate shares x's type. - // a and inv_b bind as device const float * in the kernel, so they stay F32. - const bool types_ok = - (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && - (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && - (mul0->type == x->type) && (sin_node->type == x->type) && - (sqr->type == x->type) && (mul1->type == x->type) && - (add->type == x->type); - // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D - const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; - const bool dim_ok = - (x->ne[2] == 1) && (x->ne[3] == 1) && - (add->ne[2] == 1) && (add->ne[3] == 1) && - (a->ne[2] == 1) && (a->ne[3] == 1) && - (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); - // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous - const bool contig_ok = - ggml_is_contiguous(x) && ggml_is_contiguous(add) && - ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); - - return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; -} - int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { - if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) { - return ggml_metal_op_snake_fused(ctx, idx); + int n_fuse = 1; + const ggml_metal_fusion * fusion = nullptr; + + if (ctx->use_fusion()) { + int n = 1; + fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n); + n_fuse = n; + + // snake activation autofuse: mul -> sin -> sqr -> mul -> add + if (fusion && fusion->id == GGML_METAL_FUSION_SNAKE) { + ctx->count_fusions(fusion); + return ggml_metal_op_snake_fused(ctx, idx); + } } ggml_tensor * op = ctx->node(idx); @@ -3775,9 +3775,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; - const bool use_fusion = ctx->use_fusion; + const bool use_fusion = ctx->use_fusion(); - const int debug_fusion = ctx->debug_fusion; + const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); @@ -3822,57 +3822,19 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { /*.o1 =*/ { bid_src1.offs }, }; - ggml_op fops[8]; - - int n_fuse = 1; - // c[0] = add(a, b[0]) // c[1] = add(c[0], b[1]) // c[2] = add(c[1], b[2]) // ... - if (use_fusion) { - fops[0] = GGML_OP_ADD; - fops[1] = GGML_OP_ADD; - fops[2] = GGML_OP_ADD; - fops[3] = GGML_OP_ADD; - fops[4] = GGML_OP_ADD; - fops[5] = GGML_OP_ADD; - fops[6] = GGML_OP_ADD; - fops[7] = GGML_OP_ADD; - - // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops - // across splits. idx_end indicates the last node in the current split - for (n_fuse = 0; n_fuse <= 6; ++n_fuse) { - if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { - break; - } - - ggml_tensor * f0 = ctx->node(idx + n_fuse); - ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); - - if (f0 != f1->src[0]) { - break; - } - - // b[0] === b[1] === ... - if (!ggml_are_same_layout(f0->src[1], f1->src[1])) { - break; - } - - // only fuse ops if src1 is in the same Metal buffer - ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]); - if (bid_fuse.metal != bid_src1.metal) { - break; - } - - //ctx->fuse_cnt[ops[n_fuse + 1]->op]++; - - args.o1[n_fuse + 1] = bid_fuse.offs; + if (use_fusion && fusion && fusion->id == GGML_METAL_FUSION_ADD_CHAIN) { + // the offsets of the fused addends are relative to the start of the src1 buffer + for (int i = 1; i < n_fuse; i++) { + args.o1[i] = ggml_metal_get_buffer_id(ctx->node(idx + i)->src[1]).offs; } - ++n_fuse; + ctx->count_fusions(fusion); - if (debug_fusion > 1 && n_fuse > 1) { + if (debug_fusion > 1) { GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse); } } @@ -4080,9 +4042,9 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; - const bool use_fusion = ctx->use_fusion; + const bool use_fusion = ctx->use_fusion(); - const int debug_fusion = ctx->debug_fusion; + const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); @@ -4110,8 +4072,6 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { /*.nbf3 =*/ { nb03 }, }; - ggml_op fops[8]; - int n_fuse = 1; ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 }; @@ -4120,55 +4080,35 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { // d[1] = mul(d[0], b) // d[2] = add(d[1], c) if (use_fusion) { - fops[0] = op->op; - fops[1] = GGML_OP_MUL; - fops[2] = GGML_OP_ADD; + int n = 1; + const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n); - for (n_fuse = 0; n_fuse <= 1; ++n_fuse) { - if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { - break; + if (fusion && (fusion->id == GGML_METAL_FUSION_NORM_MUL || fusion->id == GGML_METAL_FUSION_NORM_MUL_ADD)) { + n_fuse = n; + + ctx->count_fusions(fusion); + + for (int i = 1; i < n_fuse; i++) { + const ggml_tensor * fn = ctx->node(idx + i); + + bid_fuse[i - 1] = ggml_metal_get_buffer_id(fn->src[1]); + + args.nef1[i] = fn->src[1]->ne[1]; + args.nef2[i] = fn->src[1]->ne[2]; + args.nef3[i] = fn->src[1]->ne[3]; + + args.nbf1[i] = fn->src[1]->nb[1]; + args.nbf2[i] = fn->src[1]->nb[2]; + args.nbf3[i] = fn->src[1]->nb[3]; } - ggml_tensor * f0 = ctx->node(idx + n_fuse); - ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); - - if (f0 != f1->src[0]) { - break; - } - - if (f1->src[1]->ne[0] != op->ne[0]) { - break; - } - - if (!ggml_is_contiguous_rows(f1->src[1])) { - break; - } - - if (f1->type != GGML_TYPE_F32) { - break; - } - - //ctx->fuse_cnt[f1->op]++; - - bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]); - - args.nef1[n_fuse + 1] = f1->src[1]->ne[1]; - args.nef2[n_fuse + 1] = f1->src[1]->ne[2]; - args.nef3[n_fuse + 1] = f1->src[1]->ne[3]; - - args.nbf1[n_fuse + 1] = f1->src[1]->nb[1]; - args.nbf2[n_fuse + 1] = f1->src[1]->nb[2]; - args.nbf3[n_fuse + 1] = f1->src[1]->nb[3]; - } - - ++n_fuse; - - if (debug_fusion > 1 && n_fuse > 1) { - if (n_fuse == 2) { - GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); - } - if (n_fuse == 3) { - GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); + if (debug_fusion > 1) { + if (n_fuse == 2) { + GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); + } + if (n_fuse == 3) { + GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); + } } } } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index f8fe50b468..4dd8ce7af6 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -8,17 +8,18 @@ extern "C" { typedef struct ggml_metal_op * ggml_metal_op_t; +struct ggml_metal_fusion; // forward decl (ggml-metal-device.h) + ggml_metal_op_t ggml_metal_op_init( ggml_metal_device_t dev, ggml_metal_cmd_buf_t cmd_buf, struct ggml_cgraph * gf, + struct ggml_metal_fusion_info * finfo, int idx_start, int idx_end, - bool use_fusion, bool use_concurrency, bool use_capture, - int debug_graph, - int debug_fusion); + int debug_graph); void ggml_metal_op_free(ggml_metal_op_t ctx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 3bd6abd06f..4cbec8645a 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -4,6 +4,7 @@ #include "ggml-backend-impl.h" #include "ggml-metal-device.h" +#include "ggml-metal-fusion.h" #include "ggml-metal-context.h" #include "ggml-metal-ops.h" #include "ggml-metal-tuning.h" @@ -906,6 +907,30 @@ static const char * ggml_backend_metal_tuning_device_token(ggml_backend_dev_t de return ggml_metal_device_id_token(ggml_metal_device_get_props(ctx_dev)->device_id); } +// generic fusion debugging API (ad-hoc proc-address mechanism): the test resolves the device +// fusion context once and passes that opaque handle to the rest of the functions +typedef void * ggml_backend_fusion_t; + +static ggml_backend_fusion_t ggml_backend_metal_fusion_get(ggml_backend_dev_t dev) { + return ggml_metal_device_get_fusion_info((ggml_metal_device_t)dev->context); +} + +static void ggml_backend_metal_fusion_stats_init(ggml_backend_fusion_t finfo) { + ggml_metal_fusion_info_stats_init((struct ggml_metal_fusion_info *) finfo); +} + +static void ggml_backend_metal_fusion_stats_reset(ggml_backend_fusion_t finfo) { + ggml_metal_fusion_info_stats_reset((struct ggml_metal_fusion_info *) finfo); +} + +static int ggml_backend_metal_fusion_stats_get(ggml_backend_fusion_t finfo, const char ** labels, uint64_t * counts, int n) { + return ggml_metal_fusion_info_stats_get((struct ggml_metal_fusion_info *) finfo, labels, counts, n); +} + +static void ggml_backend_metal_fusion_set_enabled(ggml_backend_fusion_t finfo, bool enabled) { + ggml_metal_fusion_info_set_enabled((struct ggml_metal_fusion_info *) finfo, enabled); +} + static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const char * name) { if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_metal_get_features; @@ -928,6 +953,23 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const if (strcmp(name, "ggml_backend_metal_tuning_device_token") == 0) { return (void *)ggml_backend_metal_tuning_device_token; } + // generic fusion debugging API (ad-hoc proc-address mechanism, not part of the official + // ggml backend interface yet; a backend that adopts it exports these exact names) + if (strcmp(name, "ggml_backend_fusion_get") == 0) { + return (void *)ggml_backend_metal_fusion_get; + } + if (strcmp(name, "ggml_backend_fusion_stats_init") == 0) { + return (void *)ggml_backend_metal_fusion_stats_init; + } + if (strcmp(name, "ggml_backend_fusion_stats_reset") == 0) { + return (void *)ggml_backend_metal_fusion_stats_reset; + } + if (strcmp(name, "ggml_backend_fusion_stats_get") == 0) { + return (void *)ggml_backend_metal_fusion_stats_get; + } + if (strcmp(name, "ggml_backend_fusion_set_enabled") == 0) { + return (void *)ggml_backend_metal_fusion_set_enabled; + } return NULL; diff --git a/ggml/src/ggml-metal/kernels/gated_delta_net.metal b/ggml/src/ggml-metal/kernels/gated_delta_net.metal index 8422d8e29f..5e4861ece3 100644 --- a/ggml/src/ggml-metal/kernels/gated_delta_net.metal +++ b/ggml/src/ggml-metal/kernels/gated_delta_net.metal @@ -15,6 +15,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, + device char * dst_fuse, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -65,6 +66,12 @@ kernel void kernel_gated_delta_net_impl( // per-(seq,head) offset within a slot const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v; + // when fused with the cache cpy, write the snapshots straight into the cache buffer using + // the slot stride; otherwise append them after the attn scores (nb_out == 0) + const bool fused = args.nb_out > 0; + const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + attn_size; + const uint slot_stride = fused ? (uint)args.nb_out : state_size_per_snap; + for (short t = 0; t < args.ne22; t++) { float s_k = 0.0f; @@ -116,7 +123,7 @@ kernel void kernel_gated_delta_net_impl( if (K > 1) { const int target_slot = (int)args.ne22 - 1 - (int)t; if (target_slot >= 0 && target_slot < (int)K) { - device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base; + device float * dst_state = (device float *)state_out + (uint)target_slot * slot_stride + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -126,7 +133,7 @@ kernel void kernel_gated_delta_net_impl( } if (K == 1) { - device float * dst_state = (device float *) (dst) + attn_size + state_out_base; + device float * dst_state = (device float *)state_out + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; @@ -158,6 +165,7 @@ kernel void kernel_gated_delta_net_impl( device const char * b, device const char * s, device char * dst, + device char * dst_fuse, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { @@ -230,7 +238,13 @@ kernel void kernel_gated_delta_net_impl( dst_attn += args.ne21*S_v; } - device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20; + // when fused with the cache cpy, write the snapshots straight into the cache buffer using + // the slot stride; otherwise append them after the attn scores (nb_out == 0) + const bool fused = args.nb_out > 0; + const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v; + const uint slot_stride = fused ? (uint)args.nb_out : S_v*S_v; + + device float * dst_state = (device float *)state_out + (i23*args.ne21 + i21)*slot_stride + i20; device T * dstt_state = (device T *) (dst_state); FOR_UNROLL (short j = 0; j < NSG; j++) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 21501574a9..6334f3ccab 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -666,7 +666,9 @@ void llama_context::sched_reserve() { // need to implement a more robust mechanism that tries a few different inputs and analyzes the results ggml_cgraph * gf = nullptr; switch (model.arch) { + case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_MINIMAX_01: + // [TAG_RESERVE_DIAG_DECAY] // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1` gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc); diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index 361114acc3..9fa2e8fc01 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -229,6 +229,7 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_set_input(inp->inp_k_decay); cb(inp->inp_k_decay, "k_decay_exp", -1); + // [TAG_RESERVE_DIAG_DECAY] inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); ggml_set_input(inp->inp_diag_decay); cb(inp->inp_diag_decay, "diag_decay_exp", -1); diff --git a/src/models/plamo2.cpp b/src/models/plamo2.cpp index d946b3cff6..ba1cea1465 100644 --- a/src/models/plamo2.cpp +++ b/src/models/plamo2.cpp @@ -142,6 +142,11 @@ llama_model_plamo2::graph::graph(const llama_model & model, const llm_graph_para cur = build_plamo2_attn_layer(inp_hybrid->get_attn(), inp_pos, cur, model, il); } + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + residual = ggml_get_rows(ctx0, residual, inp_out_ids); + } + // post_mixer_norm cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il); cb(cur, "attn_post_norm", il); @@ -167,11 +172,6 @@ llama_model_plamo2::graph::graph(const llama_model & model, const llm_graph_para cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_post_norm", il); - if (il == n_layer - 1 && inp_out_ids) { - cur = ggml_get_rows(ctx0, cur, inp_out_ids); - residual = ggml_get_rows(ctx0, residual, inp_out_ids); - } - // residual connection cur = ggml_add(ctx0, cur, residual); cb(cur, "ffn_residual", il); diff --git a/src/models/qwen3vl.cpp b/src/models/qwen3vl.cpp index 5596620f07..30c08ed35e 100644 --- a/src/models/qwen3vl.cpp +++ b/src/models/qwen3vl.cpp @@ -18,6 +18,7 @@ void llama_model_qwen3vl::load_arch_tensors(llama_model_loader &) { int64_t n_vocab_out = n_vocab; if (arch == LLM_ARCH_QWEN3TTS) { + // [TAG_LLAMA_N_VOCAB_OUT] n_vocab_out = 3072; } diff --git a/tests/.gitignore b/tests/.gitignore index 52b292b1f8..04095c9ddb 100644 --- a/tests/.gitignore +++ b/tests/.gitignore @@ -1,6 +1,7 @@ * !*.* !snapshots/ +!fusion/ *.o ggml-common.h **/*.swp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5531c4ce3c..920c58c738 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,7 +196,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW - llama_build_and_test(test-llama-archs.cpp) + llama_build(test-llama-archs.cpp) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") @@ -255,6 +255,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) ARGS --models "${MODEL_DIR}" ) set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models) + + llama_build(test-fusion.cpp) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv new file mode 100644 index 0000000000..067316abfe --- /dev/null +++ b/tests/fusion/MTL.csv @@ -0,0 +1,154 @@ +# test-fusion baseline for device MTL +# arch ,moe ,mode ,label , count +arcee ,0 ,any ,RMS_NORM+MUL , 5 +arctic ,0 ,any ,RMS_NORM+MUL , 7 +baichuan ,0 ,any ,RMS_NORM+MUL , 5 +bailingmoe ,1 ,any ,ADD+ADD , 2 +bailingmoe ,1 ,any ,RMS_NORM+MUL , 5 +bailingmoe2 ,1 ,any ,ADD+ADD , 1 +bailingmoe2 ,1 ,any ,RMS_NORM+MUL , 9 +bailingmoe3 ,1 ,any ,ADD+ADD , 1 +bailingmoe3 ,1 ,any ,GATED_DELTA_NET+CPY , 1 +bailingmoe3 ,1 ,any ,RMS_NORM+MUL , 8 +bloom ,0 ,any ,NORM+MUL+ADD , 6 +chatglm ,0 ,any ,RMS_NORM+MUL , 5 +codeshell ,0 ,any ,NORM+MUL+ADD , 5 +cogvlm ,0 ,any ,RMS_NORM+MUL , 5 +command-r ,0 ,any ,NORM+MUL , 3 +dbrx ,0 ,any ,NORM+MUL , 5 +deci ,0 ,any ,RMS_NORM+MUL , 5 +deepseek ,0 ,any ,ADD+ADD , 1 +deepseek ,0 ,any ,RMS_NORM+MUL , 5 +deepseek2 ,0 ,any ,ADD+ADD , 1 +deepseek2 ,0 ,any ,RMS_NORM+MUL , 9 +deepseek32 ,0 ,any ,ADD+ADD , 1 +deepseek32 ,0 ,any ,NORM+MUL+ADD , 2 +deepseek32 ,0 ,any ,RMS_NORM+MUL , 9 +deepseek4 ,0 ,any ,RMS_NORM+MUL , 20 +dots1 ,0 ,any ,ADD+ADD , 1 +dots1 ,0 ,any ,RMS_NORM+MUL , 9 +dream ,0 ,any ,RMS_NORM+MUL , 5 +ernie4_5-moe ,1 ,any ,ADD+ADD , 1 +ernie4_5-moe ,1 ,any ,RMS_NORM+MUL , 5 +ernie4_5 ,0 ,any ,RMS_NORM+MUL , 5 +exaone ,0 ,any ,RMS_NORM+MUL , 5 +exaone4 ,0 ,any ,RMS_NORM+MUL , 5 +exaone4 ,0 ,any ,RMS_NORM+MUL+ADD , 4 +falcon ,0 ,any ,ADD+ADD , 2 +falcon ,0 ,any ,NORM+MUL+ADD , 5 +falcon-h1 ,0 ,any ,ADD+ADD , 2 +falcon-h1 ,0 ,any ,RMS_NORM+MUL , 9 +gemma ,0 ,any ,RMS_NORM+MUL , 5 +gemma2 ,0 ,any ,RMS_NORM+MUL , 5 +gemma2 ,0 ,any ,RMS_NORM+MUL+ADD , 4 +glm-dsa ,0 ,any ,ADD+ADD , 1 +glm-dsa ,0 ,any ,NORM+MUL+ADD , 2 +glm-dsa ,0 ,any ,RMS_NORM+MUL , 9 +glm4 ,0 ,any ,RMS_NORM+MUL , 5 +glm4 ,0 ,any ,RMS_NORM+MUL+ADD , 4 +glm4moe ,1 ,any ,ADD+ADD , 1 +glm4moe ,1 ,any ,RMS_NORM+MUL , 9 +gpt-oss ,0 ,any ,RMS_NORM+MUL , 5 +gpt2 ,0 ,any ,NORM+MUL+ADD , 5 +gptneox ,0 ,any ,NORM+MUL+ADD , 5 +granite ,0 ,any ,RMS_NORM+MUL , 5 +granite ,0 ,any ,RMS_NORM+MUL , 5 +granitehybrid ,0 ,any ,RMS_NORM+MUL , 6 +granitemoe ,1 ,any ,RMS_NORM+MUL , 5 +granitemoe ,1 ,any ,RMS_NORM+MUL , 5 +grok ,0 ,any ,RMS_NORM+MUL , 5 +grok ,0 ,any ,RMS_NORM+MUL+ADD , 4 +grovemoe ,1 ,any ,ADD+ADD , 2 +grovemoe ,1 ,any ,RMS_NORM+MUL , 9 +hunyuan-dense ,0 ,any ,RMS_NORM+MUL , 9 +hunyuan-moe ,1 ,any ,ADD+ADD , 2 +hunyuan-moe ,1 ,any ,RMS_NORM+MUL , 9 +hunyuan_vl ,0 ,any ,RMS_NORM+MUL , 9 +hy_v3 ,0 ,any ,ADD+ADD , 2 +hy_v3 ,0 ,any ,RMS_NORM+MUL , 9 +hy_v4 ,0 ,any ,NORM+MUL+ADD , 1 +hy_v4 ,0 ,any ,RMS_NORM+MUL , 9 +internlm2 ,0 ,any ,RMS_NORM+MUL , 5 +jais ,0 ,any ,NORM+MUL+ADD , 5 +jais2 ,0 ,any ,NORM+MUL+ADD , 5 +jamba ,0 ,any ,RMS_NORM+MUL , 8 +kimi-k3 ,0 ,any ,GATED_DELTA_NET+CPY , 1 +kimi-k3 ,0 ,any ,RMS_NORM+MUL , 17 +kimi-linear ,0 ,any ,ADD+ADD , 1 +kimi-linear ,0 ,any ,GATED_DELTA_NET+CPY , 1 +kimi-linear ,0 ,any ,RMS_NORM+MUL , 7 +lfm2 ,0 ,any ,RMS_NORM+MUL , 7 +lfm2moe ,1 ,any ,RMS_NORM+MUL , 7 +llada ,0 ,any ,RMS_NORM+MUL , 5 +llada-moe ,1 ,any ,RMS_NORM+MUL , 9 +llama ,0 ,any ,RMS_NORM+MUL , 5 +llama ,0 ,any ,RMS_NORM+MUL , 5 +llama4 ,0 ,any ,ADD+ADD , 2 +llama4 ,0 ,any ,RMS_NORM+MUL , 9 +maincoder ,0 ,any ,RMS_NORM+MUL , 9 +mamba ,0 ,any ,RMS_NORM+MUL , 3 +mamba2 ,0 ,any ,RMS_NORM+MUL , 5 +minicpm ,0 ,any ,RMS_NORM+MUL , 5 +minicpm ,0 ,any ,RMS_NORM+MUL , 5 +minicpm3 ,0 ,any ,RMS_NORM+MUL , 9 +minimax-01 ,0 ,any ,RMS_NORM+MUL , 6 +minimax-m2 ,0 ,any ,RMS_NORM+MUL , 9 +minimax-m3 ,0 ,any ,ADD+ADD , 1 +minimax-m3 ,0 ,any ,RMS_NORM+MUL , 11 +mistral3 ,0 ,any ,RMS_NORM+MUL , 5 +mistral3 ,0 ,any ,RMS_NORM+MUL , 5 +mistral4 ,0 ,any ,ADD+ADD , 1 +mistral4 ,0 ,any ,RMS_NORM+MUL , 9 +mpt ,0 ,any ,NORM+MUL+ADD , 5 +nanbeige ,0 ,any ,RMS_NORM+MUL , 5 +nemotron ,0 ,any ,NORM+MUL+ADD , 5 +nemotron_h ,0 ,any ,RMS_NORM+MUL , 5 +nemotron_h_moe ,1 ,any ,RMS_NORM+MUL , 5 +olmoe ,1 ,any ,RMS_NORM+MUL , 9 +openelm ,0 ,any ,RMS_NORM+MUL , 9 +orion ,0 ,any ,NORM+MUL+ADD , 5 +paddleocr ,0 ,any ,RMS_NORM+MUL , 5 +pangu-embedded ,0 ,any ,RMS_NORM+MUL , 5 +phi2 ,0 ,any ,ADD+ADD , 2 +phi2 ,0 ,any ,NORM+MUL+ADD , 3 +phi3 ,0 ,any ,RMS_NORM+MUL , 5 +phimoe ,1 ,any ,RMS_NORM+MUL+ADD , 5 +plamo ,0 ,any ,ADD+ADD , 2 +plamo ,0 ,any ,RMS_NORM+MUL , 3 +plamo2 ,0 ,any ,RMS_NORM+MUL , 10 +plamo2 ,0 ,any ,RMS_NORM+MUL+ADD , 4 +pockettts ,0 ,any ,NORM+MUL+ADD , 5 +qwen ,0 ,any ,RMS_NORM+MUL , 5 +qwen2 ,0 ,any ,RMS_NORM+MUL , 5 +qwen2moe ,1 ,any ,ADD+ADD , 2 +qwen2moe ,1 ,any ,RMS_NORM+MUL , 5 +qwen2vl ,0 ,any ,RMS_NORM+MUL , 5 +qwen3 ,0 ,any ,RMS_NORM+MUL , 9 +qwen35 ,0 ,any ,GATED_DELTA_NET+CPY , 1 +qwen35 ,0 ,any ,RMS_NORM+MUL , 8 +qwen35moe ,1 ,any ,ADD+ADD , 2 +qwen35moe ,1 ,any ,GATED_DELTA_NET+CPY , 1 +qwen35moe ,1 ,any ,RMS_NORM+MUL , 8 +qwen3moe ,1 ,any ,RMS_NORM+MUL , 9 +qwen3next ,0 ,any ,ADD+ADD , 2 +qwen3next ,0 ,any ,GATED_DELTA_NET+CPY , 1 +qwen3next ,0 ,any ,RMS_NORM+MUL , 8 +qwen3tts ,0 ,any ,RMS_NORM+MUL , 9 +qwen3vl ,0 ,any ,RMS_NORM+MUL , 9 +qwen3vlmoe ,1 ,any ,RMS_NORM+MUL , 9 +qwen4exp ,0 ,any ,ADD+ADD+ADD , 5 +qwen4exp ,0 ,any ,ADD+ADD+ADD+ADD+ADD+ADD+ADD , 9 +qwen4exp ,0 ,any ,GATED_DELTA_NET+CPY , 1 +qwen4exp ,0 ,any ,RMS_NORM+MUL , 5 +refact ,0 ,any ,RMS_NORM+MUL , 5 +refact ,0 ,any ,RMS_NORM+MUL , 5 +rnd1 ,0 ,any ,RMS_NORM+MUL , 9 +seed_oss ,0 ,any ,RMS_NORM+MUL , 5 +smallthinker ,0 ,any ,RMS_NORM+MUL , 5 +smollm3 ,0 ,any ,RMS_NORM+MUL , 5 +stablelm ,0 ,any ,NORM+MUL , 4 +stablelm ,0 ,any ,NORM+MUL+ADD , 5 +starcoder ,0 ,any ,NORM+MUL+ADD , 5 +starcoder2 ,0 ,any ,NORM+MUL+ADD , 5 +talkie ,0 ,any ,ADD+ADD , 2 +xverse ,0 ,any ,RMS_NORM+MUL , 5 diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ef4fc30cec..15c42a1081 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4638,6 +4638,122 @@ struct test_gated_delta_net : public test_case { } }; +// GGML_OP_GATED_DELTA_NET + GGML_OP_CPY (recurrent cache fusion) +struct test_gated_delta_net_cache_fusion : public test_case { + const ggml_type type; + + const int64_t head_count; + const int64_t head_size; + const int64_t n_seq_tokens; + const int64_t n_seqs; + const int64_t K; // snapshot slot count (>1) + + ggml_tensor * cpy_node = nullptr; + + std::string vars() override { + return VARS_TO_STR6(type, head_count, head_size, n_seq_tokens, n_seqs, K); + } + + test_gated_delta_net_cache_fusion(ggml_type type = GGML_TYPE_F32, + int64_t head_count = 4, int64_t head_size = 32, int64_t n_seq_tokens = 2, int64_t n_seqs = 1, + int64_t K = 2) + : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + const int64_t S_v = head_size; + const int64_t H_v = head_count; + const int64_t H_k = head_count; + const int64_t D = S_v * S_v * H_v; + const int64_t n_written = std::min(n_seq_tokens, K); + + ggml_tensor * q = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs); + ggml_tensor * v = ggml_new_tensor_4d(ctx, type, head_size, H_v, n_seq_tokens, n_seqs); + ggml_set_name(q, "q"); + ggml_set_name(k, "k"); + ggml_set_name(v, "v"); + ggml_tensor * g = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs); + ggml_tensor * beta = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs); + ggml_tensor * state = ggml_new_tensor_4d(ctx, type, head_size, head_size, H_v, n_seqs); + ggml_set_name(g, "g"); + ggml_set_name(beta, "beta"); + ggml_set_name(state, "state"); + + q = ggml_l2_norm(ctx, q, 1e-6f); + k = ggml_l2_norm(ctx, k, 1e-6f); + + ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, K); + ggml_set_name(gdn_out, "gdn_out"); + + // attn scores view (first part of the gdn output) + ggml_tensor * attn = ggml_view_4d(ctx, gdn_out, + S_v, H_v, n_seq_tokens, n_seqs, + ggml_row_size(gdn_out->type, S_v), + ggml_row_size(gdn_out->type, S_v * H_v), + ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens), 0); + ggml_set_name(attn, "attn"); + + // snapshot tail view [D, n_seqs, n_written] + const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs; + ggml_tensor * src = ggml_view_3d(ctx, gdn_out, + D, n_seqs, n_written, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, D * n_seqs), + ggml_row_size(gdn_out->type, attn_score_elems)); + + // recurrent cache view [D, n_seqs, n_written] + ggml_tensor * cache = ggml_new_tensor_3d(ctx, type, D, n_seqs, n_written); + ggml_set_name(cache, "cache"); + ggml_tensor * dst = ggml_view_3d(ctx, cache, + D, n_seqs, n_written, + ggml_row_size(cache->type, D), + ggml_row_size(cache->type, D * n_seqs), 0); + + ggml_tensor * cpy = ggml_cpy(ctx, src, dst); + ggml_set_name(cpy, "gdn_cache_cpy"); + cpy_node = cpy; + + // read the cpy output (not the plain dst view, which would not pull the cpy into the graph) + // so that neither the gdn nor the cpy is the graph output + ggml_tensor * out = ggml_sum(ctx, cpy); + return out; + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "GATED_DELTA_NET_CACHE_FUSION"; + } + + bool run_whole_graph() override { return true; } + std::vector fusion_test_nodes() override { return { cpy_node }; } + + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + const uint64_t S_v = head_size; + const uint64_t H_v = head_count; + const uint64_t T = n_seq_tokens; + const uint64_t B = n_seqs; + return (4ull*S_v + 2ull*S_v*S_v) * H_v * T * B; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (ggml_is_view_op(t->op)) { continue; } + if (strcmp(t->name, "g") == 0) { + init_tensor_uniform(t, -20.0f, -1e-4f); + } else if (strcmp(t->name, "beta") == 0) { + init_tensor_uniform(t, 0.0f, 1.0f); + } else if (strcmp(t->name, "v") == 0) { + init_tensor_uniform(t, -0.3f, 5.0f); + } else if (strcmp(t->name, "cache") == 0) { + init_tensor_uniform(t, 0.0f, 0.0f); + } else { + init_tensor_uniform(t); + } + } + } +}; + // GGML_OP_GATED_LINEAR_ATTN struct test_gla : public test_case { const ggml_type type; @@ -10741,6 +10857,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4)); + // gdn + cache cpy fusion (K > 1) + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 2, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 64, 4, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 4, 1, 4)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 8, 32, 4, 2, 4)); + test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 8, 1, 4)); + #if 0 // these tests are disabled to save execution time, sbut they can be handy for debugging test_cases.emplace_back(new test_llama(2, true)); diff --git a/tests/test-fusion.cpp b/tests/test-fusion.cpp new file mode 100644 index 0000000000..467248f0f1 --- /dev/null +++ b/tests/test-fusion.cpp @@ -0,0 +1,565 @@ +// test-fusion: verify the backend fusion logic against a per-device baseline. +// +// for every dummy model generated by test-llama-archs, the tool runs the model on a single +// device with fusion enabled and disabled, and reports: +// - the per-fusion-type counters for each mode (prefill / decode, merged into "any" when the +// per-graph counts match) +// - the NMSE between the fused and unfused logits +// - the NMSE between the device and a CPU reference +// +// the per-fusion-type counters are compared against a per-device baseline file (CSV) so a +// fusion pattern that silently stops matching (or fires when it should not) is caught as a +// regression. +// +// usage: +// test-fusion --models DIR --device MTL0 --record baseline.csv # generate a baseline +// test-fusion --models DIR --device MTL0 --check baseline.csv # validate against it +// test-fusion --model FILE --device MTL0 --check baseline.csv # validate a single model + +#include "common.h" +#include "log.h" +#include "llama-cpp.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// generic fusion debugging API, resolved through the ad-hoc get_proc_address mechanism +// (not part of the official ggml backend interface yet). a backend that adopts fusion debugging +// exports these exact names. +typedef void * ggml_backend_fusion_t; + +typedef ggml_backend_fusion_t ( * fusion_get_t) (ggml_backend_dev_t); +typedef void ( * fusion_stats_init_t) (ggml_backend_fusion_t); +typedef void ( * fusion_stats_reset_t) (ggml_backend_fusion_t); +typedef int ( * fusion_stats_get_t) (ggml_backend_fusion_t, const char **, uint64_t *, int); +typedef void ( * fusion_set_enabled_t) (ggml_backend_fusion_t, bool); + +static bool silent_model_load_progress(float, void *) { + return true; +} + +struct gguf_context_ptr { + gguf_context * ctx; + gguf_context_ptr(gguf_context * c) : ctx(c) {} + ~gguf_context_ptr() { if (ctx) { gguf_free(ctx); } } + gguf_context * get() const { return ctx; } + gguf_context_ptr(const gguf_context_ptr &) = delete; + gguf_context_ptr & operator=(const gguf_context_ptr &) = delete; +}; + +// NMSE between two vectors (same as tests/test-llama-archs.cpp) +static double nmse(const std::vector & a, const std::vector & b) { + GGML_ASSERT(a.size() == b.size()); + double mse_a_b = 0.0; + double mse_a_0 = 0.0; + + for (size_t i = 0; i < a.size(); i++) { + const float a_i = a[i]; + const float b_i = b[i]; + + mse_a_b += (a_i - b_i) * (a_i - b_i); + mse_a_0 += a_i * a_i; + } + + return mse_a_b / mse_a_0; +} + +// deterministic token sequence +static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed) { + std::mt19937 gen(seed); + std::uniform_int_distribution<> dis(0, n_vocab - 1); + std::vector ret; + ret.reserve(n_tokens); + for (uint32_t i = 0; i < n_tokens; i++) { + ret.push_back(dis(gen)); + } + return ret; +} + +// trim leading/trailing whitespace (used when parsing padded CSV columns) +static std::string trim(const std::string & s) { + const size_t b = s.find_first_not_of(" \t\r\n"); + if (b == std::string::npos) { + return ""; + } + const size_t e = s.find_last_not_of(" \t\r\n"); + return s.substr(b, e - b + 1); +} + +static std::string get_arch(const std::string & path) { + gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/nullptr }; + gguf_context_ptr ctx(gguf_init_from_file(path.c_str(), params)); + if (!ctx.get()) { + throw std::runtime_error("failed to read gguf: " + path); + } + const int idx = gguf_find_key(ctx.get(), "general.architecture"); + if (idx < 0) { + return "unknown"; + } + const char * val = gguf_get_val_str(ctx.get(), idx); + return val ? val : "unknown"; +} + +static llama_model_ptr load_model(const std::string & path, ggml_backend_dev_t dev) { + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + std::vector devs = { dev, nullptr }; + model_params.devices = devs.data(); + model_params.split_mode = LLAMA_SPLIT_MODE_LAYER; + + llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params)); + if (!model) { + throw std::runtime_error("failed to load model: " + path); + } + return model; +} + +// a fresh context (fresh state) from an already-loaded model +static llama_context_ptr create_ctx(llama_model * model, int n_ubatch) { + llama_context_params ctx_params = llama_context_default_params(); + ctx_params.n_ctx = 0; + ctx_params.n_threads = 4; + ctx_params.n_threads_batch = 4; + ctx_params.n_ubatch = n_ubatch; + ctx_params.n_batch = n_ubatch; + + llama_context_ptr lctx(llama_init_from_model(model, ctx_params)); + if (!lctx) { + throw std::runtime_error("failed to init context"); + } + return lctx; +} + +// decode all tokens in one batch; returns the logits of every token +static std::vector decode_prefill(llama_model * model, llama_context * lctx, const std::vector & tokens) { + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + llama_batch batch = llama_batch_init(tokens.size(), 0, 1); + for (size_t i = 0; i < tokens.size(); i++) { + common_batch_add(batch, tokens[i], i, { 0 }, true); + } + batch.n_tokens = tokens.size(); + if (llama_decode(lctx, batch)) { + llama_batch_free(batch); + throw std::runtime_error("prefill decode failed"); + } + + std::vector ret; + ret.reserve(tokens.size() * n_vocab); + for (size_t i = 0; i < tokens.size(); i++) { + const float * logits_ith = llama_get_logits_ith(lctx, i); + for (uint32_t j = 0; j < n_vocab; j++) { + ret.push_back(logits_ith[j]); + } + } + llama_batch_free(batch); + return ret; +} + +// decode one token at a time; returns the logits of the last token of each step +static std::vector decode_gen(llama_model * model, llama_context * lctx, const std::vector & tokens) { + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + llama_batch batch = llama_batch_init(1, 0, 1); + std::vector ret; + for (size_t i = 0; i < tokens.size(); i++) { + common_batch_clear(batch); + common_batch_add(batch, tokens[i], i, { 0 }, true); + if (llama_decode(lctx, batch)) { + llama_batch_free(batch); + throw std::runtime_error("decode failed"); + } + const float * logits = llama_get_logits_ith(lctx, 0); + for (uint32_t j = 0; j < n_vocab; j++) { + ret.push_back(logits[j]); + } + } + llama_batch_free(batch); + return ret; +} + +static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t finfo, + std::vector & labels, std::vector & counts) { + const int n = api_stats_get(finfo, nullptr, nullptr, 0); + labels.assign(n, nullptr); + counts.assign(n, 0); + api_stats_get(finfo, labels.data(), counts.data(), n); +} + +// one row of the per-label report +struct fusion_row { + std::string arch; + bool moe; + std::string mode; + std::string label; + uint64_t count_fused; + uint64_t count_unfused; + uint64_t expected; + double nmse_fus; + double nmse_dev; + bool ok_count; // counts match the baseline + bool ok_nmse; // nmse within epsilon +}; + +static void usage(const char * argv0) { + printf("%s: verify fusion counts on a device against a per-device baseline\n\n", argv0); + printf("usage: %s [options]\n\n", argv0); + printf("options:\n"); + printf(" --models DIR run over all .gguf models in a directory\n"); + printf(" --model FILE run over a single model file (mutually exclusive with --models)\n"); + printf(" --device NAME device to run on (e.g. MTL0, CPU)\n"); + printf(" --record CSV write the golden baseline\n"); + printf(" --check CSV validate the counters against a baseline (default)\n"); + printf(" -h, --help show this message and exit\n"); +} + +int main(int argc, char ** argv) { + std::string models_dir; + std::string model_file; + std::string device_name; + std::string record_path; + std::string check_path; + + for (int i = 1; i < argc; i++) { + const std::string arg = argv[i]; + const auto next = [&](const char * name) -> std::string { + if (i + 1 >= argc) { + LOG_ERR("%s: %s requires an argument\n", __func__, name); + exit(1); + } + return argv[++i]; + }; + if (arg == "-h" || arg == "--help") { + usage(argv[0]); + exit(0); + } + if (arg == "--models") { models_dir = next("--models"); } + else if (arg == "--model") { model_file = next("--model"); } + else if (arg == "--device"){ device_name = next("--device"); } + else if (arg == "--record"){ record_path = next("--record"); } + else if (arg == "--check") { check_path = next("--check"); } + else { + LOG_ERR("%s: unknown argument: %s\n", __func__, arg.c_str()); + return 1; + } + } + + if (device_name.empty()) { + LOG_ERR("%s: --device NAME is required\n", __func__); + return 1; + } + if (models_dir.empty() && model_file.empty()) { + LOG_ERR("%s: --models DIR or --model FILE is required\n", __func__); + return 1; + } + if (!models_dir.empty() && !model_file.empty()) { + LOG_ERR("%s: --models DIR and --model FILE are mutually exclusive\n", __func__); + return 1; + } + if (!record_path.empty() && !check_path.empty()) { + LOG_ERR("%s: --record and --check are mutually exclusive\n", __func__); + return 1; + } + + std::vector models; + if (!model_file.empty()) { + if (!std::filesystem::is_regular_file(model_file)) { + LOG_ERR("%s: model file '%s' does not exist\n", __func__, model_file.c_str()); + return 1; + } + models.push_back(model_file); + } else { + if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { + LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str()); + return 1; + } + for (const auto & entry : std::filesystem::directory_iterator(models_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".gguf") { + models.push_back(entry.path().string()); + } + } + std::sort(models.begin(), models.end()); + + if (models.empty()) { + LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str()); + return 1; + } + } + + common_init(); + ggml_backend_load_all(); + + ggml_backend_dev_t dev = ggml_backend_dev_by_name(device_name.c_str()); + if (!dev) { + LOG_WRN("%s: device '%s' not found - skipping (baseline is device-specific)\n", + __func__, device_name.c_str()); + return 0; + } + + // resolve the generic fusion debugging functions through the ad-hoc get_proc_address + // mechanism; a backend that does not adopt fusion debugging exports none of them + auto * reg = ggml_backend_dev_backend_reg(dev); + + // output naming uses the backend base name (e.g. "MTL") rather than the specific device + // name (e.g. "MTL0") the test was invoked with + const std::string base_name = ggml_backend_reg_name(reg); + + auto api_get = (fusion_get_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_get"); + auto api_stats_init = (fusion_stats_init_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_init"); + auto api_stats_reset = (fusion_stats_reset_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_reset"); + auto api_stats_get = (fusion_stats_get_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_get"); + auto api_set_enabled = (fusion_set_enabled_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_set_enabled"); + + if (!api_get || !api_stats_init || !api_set_enabled || !api_stats_reset || !api_stats_get) { + LOG_ERR("%s: device '%s' does not export the generic fusion debugging API " + "(ggml_backend_fusion_*) - cannot run the fusion regression test\n", + __func__, device_name.c_str()); + return 1; + } + + ggml_backend_fusion_t finfo = api_get(dev); + + // enable fusions stats + api_stats_init(finfo); + + const bool has_counts = true; + + // load the baseline (if any): key arch|moe|mode|label -> expected count + std::map baseline; + if (!check_path.empty()) { + std::ifstream in(check_path); + if (!in) { + LOG_ERR("%s: cannot open baseline '%s'\n", __func__, check_path.c_str()); + return 1; + } + std::string line; + while (std::getline(in, line)) { + if (line.empty() || line[0] == '#') { + continue; + } + std::vector cols; + size_t pos = 0; + while ((pos = line.find(',')) != std::string::npos) { + cols.push_back(trim(line.substr(0, pos))); + line.erase(0, pos + 1); + } + cols.push_back(trim(line)); + if (cols.size() != 5) { + continue; + } + baseline[cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3]] = std::stoull(cols[4]); + } + } + + std::vector rows; + + LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str()); + + const size_t seed = 1; + + for (const auto & model_path : models) { + const std::string arch = get_arch(model_path); + const bool moe = arch.find("moe") != std::string::npos; + + llama_model_ptr model; + llama_model_ptr model_cpu; + uint32_t n_vocab = 0; + try { + model = load_model(model_path, dev); + model_cpu = load_model(model_path, ggml_backend_dev_by_name("CPU")); + n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get())); + } catch (const std::exception & e) { + LOG_ERR("%s: %s: %s\n", __func__, model_path.c_str(), e.what()); + continue; + } + + struct mode_cfg { + std::string name; + std::vector (*decode)(llama_model *, llama_context *, const std::vector &); + int n_tokens; + int n_graphs; // graph runs per mode (prefill=1, decode=16) + }; + const mode_cfg modes[] = { + { "prefill", decode_prefill, 32, 1 }, + { "decode", decode_gen, 16, 16 }, + }; + + // per-label, per-mode data for this model; prefill and decode are merged into a single + // "any" row when their per-graph counts match + struct mode_data { + bool present; + uint64_t count_fused; // per graph + uint64_t count_unfused; // per graph + double nmse_fus; + double nmse_dev; + bool ok_nmse; + }; + std::map> mdata; + + for (int mi = 0; mi < 2; mi++) { + const mode_cfg & mode = modes[mi]; + const auto tokens = get_tokens(mode.n_tokens, n_vocab, seed); + + // CPU reference for this mode (fresh context, fresh state) + std::vector logits_cpu; + try { + llama_context_ptr ctx = create_ctx(model_cpu.get(), 32); + logits_cpu = mode.decode(model_cpu.get(), ctx.get(), tokens); + } catch (const std::exception & e) { + LOG_WRN("%s: %s: cpu reference: %s\n", __func__, model_path.c_str(), e.what()); + } + + // fused run on a fresh context (fresh state) + std::vector logits_fused; + std::vector labels; + std::vector counts_fused; + { + llama_context_ptr ctx = create_ctx(model.get(), 32); + if (has_counts) { + api_set_enabled(finfo, true); + api_stats_reset(finfo); + } + logits_fused = mode.decode(model.get(), ctx.get(), tokens); + if (has_counts) { + read_counts(api_stats_get, finfo, labels, counts_fused); + } + } + + // unfused run on another fresh context (fresh state) + std::vector logits_unfused; + std::vector counts_unfused; + { + llama_context_ptr ctx = create_ctx(model.get(), 32); + if (has_counts) { + api_set_enabled(finfo, false); + api_stats_reset(finfo); + } + logits_unfused = mode.decode(model.get(), ctx.get(), tokens); + if (has_counts) { + read_counts(api_stats_get, finfo, labels, counts_unfused); + } + } + + const double nmse_fus = nmse(logits_fused, logits_unfused); + const double nmse_dev = logits_cpu.empty() ? 0.0 : nmse(logits_fused, logits_cpu); + + if (has_counts) { + for (int i = 0; i < (int) labels.size(); i++) { + const uint64_t fused = counts_fused[i] / mode.n_graphs; + const uint64_t unfused = counts_unfused[i] / mode.n_graphs; + if (fused == 0 && unfused == 0) { + continue; + } + auto & d = mdata[labels[i]][mi]; + d.present = true; + d.count_fused = fused; + d.count_unfused = unfused; + d.nmse_fus = nmse_fus; + d.nmse_dev = nmse_dev; + d.ok_nmse = nmse_fus <= 1e-4; + } + } else { + rows.push_back({ arch, moe, mode.name, "?", 0, 0, 0, nmse_fus, nmse_dev, true, nmse_fus <= 1e-4 }); + } + } + + // build the per-label rows, merging prefill and decode into "any" when the per-graph + // counts match (they always do for the deterministic fusion table) + if (has_counts) { + for (auto & kv : mdata) { + const std::string & label = kv.first; + const auto & d = kv.second; + const bool both = d[0].present && d[1].present; + const bool match = both && d[0].count_fused == d[1].count_fused; + + if (match) { + // one "any" row; use the worst NMSE across the two modes + const std::string any_key = arch + "|" + (moe ? "1" : "0") + "|any|" + label; + const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key) : 0; + const bool ok_count = check_path.empty() || d[0].count_fused == expected; + const bool ok_nmse = d[0].ok_nmse && d[1].ok_nmse; + const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus); + const double nmse_dev = std::max(d[0].nmse_dev, d[1].nmse_dev); + rows.push_back({ arch, moe, "any", label, d[0].count_fused, d[0].count_unfused, + expected, nmse_fus, nmse_dev, ok_count, ok_nmse }); + } else { + // counts differ - keep a separate row per mode + for (int mi = 0; mi < 2; mi++) { + if (!d[mi].present) { + continue; + } + const mode_data & a = d[mi]; + const std::string mode_key = arch + "|" + (moe ? "1" : "0") + "|" + modes[mi].name + "|" + label; + const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key) : 0; + const bool ok_count = check_path.empty() || a.count_fused == expected; + rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused, + expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse }); + } + } + } + } + + LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str()); + } + + // print the report + { + std::ofstream out(record_path); + std::ostream & os = record_path.empty() ? std::cout : out; + if (!record_path.empty()) { + os << "# test-fusion baseline for device " << base_name << "\n"; + os << "# " << std::left + << std::setw(18) << "arch" << ',' + << std::setw(4) << "moe" << ',' + << std::setw(8) << "mode" << ',' + << std::setw(28) << "label" << ',' + << std::right << std::setw(7) << "count" << '\n'; + } + + LOG_INF("%-20s %-4s %-8s %-22s %7s %7s %7s %10s %10s %s\n", + "arch", "moe", "mode", "label", "fused", "unfused", "expected", "nmse_fus", "nmse_dev", "status"); + int n_ok = 0; + int n_bad = 0; + for (const auto & r : rows) { + const bool ok = r.ok_count && r.ok_nmse; + const char * status = ok ? "ok" : "FAIL"; + if (ok) { n_ok++; } else { n_bad++; } + LOG_INF("%-20s %-4s %-8s %-22s %7llu %7llu %7llu %10.2e %10.2e %s\n", + r.arch.c_str(), r.moe ? "moe" : "dense", r.mode.c_str(), r.label.c_str(), + (unsigned long long) r.count_fused, (unsigned long long) r.count_unfused, + (unsigned long long) r.expected, r.nmse_fus, r.nmse_dev, status); + if (!record_path.empty()) { + os << std::left + << std::setw(20) << r.arch << ',' + << std::setw(4) << (r.moe ? "1" : "0") << ',' + << std::setw(8) << r.mode << ',' + << std::setw(28) << r.label << ',' + << std::right << std::setw(7) << r.count_fused << '\n'; + } + } + LOG_INF("summary: %d ok, %d failed\n", n_ok, n_bad); + if (!record_path.empty()) { + LOG_INF("%s: baseline written to '%s'\n", __func__, record_path.c_str()); + } + + if (n_bad && !models_dir.empty() && !check_path.empty()) { + LOG_WRN("%s: if the fusion counts are expected to change, run with --record to update the baseline:\n" + "\n" + "./bin/test-llama-archs -o %s\n" + "%s --device %s --models %s --record %s\n", + __func__, models_dir.c_str(), argv[0], device_name.c_str(), models_dir.c_str(), check_path.c_str()); + } + + return n_bad; + } +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index dbed9846f9..3496f72e49 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -128,7 +128,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } else if (arch == LLM_ARCH_CHAMELEON) { n_vocab = 10240; } else if (arch == LLM_ARCH_QWEN3TTS) { - n_vocab = 4096; // must be >= the hard-coded codec head size (3072) + //n_vocab = 4096; // must be >= the hard-coded codec head size (3072) + n_vocab = 3072; // TODO: should be 4096, but user code cannot get `n_vocab_out` yet [TAG_LLAMA_N_VOCAB_OUT] } uint32_t n_head_kv = n_head; diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6179e6c108..74d1ba6c21 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -109,7 +109,7 @@ static bool test_seq_rm_isolated( for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) { llama_batch_ptr batch(n_tokens, 0, 1); for (size_t i = 0; i < n_tokens; ++i) { - common_batch_add(batch.get(), tokens[i], i, { seq_id }, false); + common_batch_add(batch.get(), tokens[i], i, { seq_id }, i == n_tokens - 1); } if (llama_decode(ctx.get(), batch.get())) { @@ -373,7 +373,7 @@ static bool test_seq_cp_scatter(struct llama_model * model, const struct common_ auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) { llama_batch_ptr batch(1, 0, 1); - common_batch_add(batch.get(), tok, pos, { seq }, false); + common_batch_add(batch.get(), tok, pos, { seq }, true); return llama_decode(ctx.get(), batch.get()) == 0; };