refactor: externalize Sol-Attn and use public graph APIs

src/core/compute_workspace.cpp

@@ -6,9 +6,9 @@
 #include <unordered_map>
 #include <unordered_set>
 
+#include "core/ggml_compat.h"
 #include "core/ggml_extend_backend.h"
 #include "core/util.h"
-#include "ggml-impl.h"
 
 namespace sd {
     ComputeWorkspace::~ComputeWorkspace() {
@@ -69,7 +69,7 @@ namespace sd {
 
     bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
         GGML_ASSERT(!active_);
-        const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+        const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
         if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
             release();
             scheduler_          = make_scheduler(capacity);
@@ -101,11 +101,11 @@ namespace sd {
                 tensors.push_back(tensor);
             }
         };
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            visit(graph->nodes[i]);
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            visit(ggml_graph_node(graph, i));
         }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            visit(graph->leafs[i]);
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            visit(ggml_graph_leaf(graph, i));
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
             visit(tensors[i]->view_src);
@@ -120,9 +120,8 @@ namespace sd {
                 const auto* bytes = reinterpret_cast<const uint8_t*>(&value);
                 key.insert(key.end(), bytes, bytes + sizeof(value));
             };
-            append(graph->n_nodes);
-            append(graph->n_leafs);
-            append(graph->order);
+            append(ggml_graph_n_nodes(graph));
+            append(ggml_graph_n_leafs(graph));
             for (auto tensor : tensors) {
                 append(tensor);
                 append(tensor->type);
@@ -151,10 +150,7 @@ namespace sd {
                     const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
                     append(alignment_offset);
                 }
-                const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
-                const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
-                                         ? graph->use_counts[index]
-                                         : 0;
+                const int32_t uses = ggml_graph_use_count(graph, tensor);
                 append(uses);
             }
             // Match the full allocation metadata, including topology and external placement.
@@ -195,24 +191,13 @@ namespace sd {
                 copy->extra  = nullptr;
             }
         }
-        auto copy_graph     = ggml_new_graph_custom(context, graph_size, false);
-        copy_graph->n_nodes = graph->n_nodes;
-        copy_graph->n_leafs = graph->n_leafs;
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            copy_graph->nodes[i] = copies.at(graph->nodes[i]);
-        }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            copy_graph->leafs[i] = copies.at(graph->leafs[i]);
-        }
-        copy_graph->order = graph->order;
-        // Preserve fusion metadata, including uses outside a graph view, without rebuilding its node order.
-        for (const auto& entry : copies) {
-            const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first);
-            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) {
-                const size_t copy_index            = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
-                copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
-            }
-        }
+        auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
+        auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* {
+            const auto& mapped = *static_cast<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);
+            const auto found   = mapped.find(tensor);
+            return found == mapped.end() ? nullptr : found->second;
+        };
+        GGML_ASSERT(ggml_graph_cpy_mapped(graph, copy_graph, map_tensor, &copies));
         Measurement result;
         auto scheduler = make_scheduler(graph_size);
         if (scheduler != nullptr) {

src/core/cuda/CMakeLists.txt

@@ -23,8 +23,8 @@ endfunction()
 if(SD_CUDA)
     enable_language(CUDA)
     sd_cuda_sources(cuda CUDA)
-    target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu)
-    set_source_files_properties(sage_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu sol_attention.cu)
+    set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
     set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
     if(GGML_CUDA_CUTLASS_PATH)
         set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)

src/core/cuda/registration.cu

@@ -13,6 +13,7 @@ namespace SD_CUDA_NAMESPACE {
     void register_affine(ggml_backend_dev_t device);
 #if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
     void register_sage_attention(ggml_backend_dev_t device);
+    void register_sol_attention(ggml_backend_dev_t device);
 #endif
 
     void register_extensions(ggml_backend_dev_t device) {
@@ -24,6 +25,7 @@ namespace SD_CUDA_NAMESPACE {
         register_rope(device);
 #if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
         register_sage_attention(device);
+        register_sol_attention(device);
 #endif
 #ifdef SD_CUDA_CUTE_ATTENTION
         sd_ggml_register_cuda_i8_projection(device);

src/core/cuda/sol_attention.cu

@@ -0,0 +1,409 @@
+// Sol-Attn diagonal-threshold algorithm:
+// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+#if __CUDACC_VER_MAJOR__ >= 12
+
+#include <cuda_bf16.h>
+#include <algorithm>
+#include <climits>
+
+namespace sd_cuda {
+
+    static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1)
+            x += __shfl_xor_sync(0xffffffff, x, offset);
+        return x;
+    }
+
+    static constexpr int SOL_DIM   = 128;
+    static constexpr int SOL_BLOCK = 64;
+
+    struct SolWorkspace {
+        size_t q, k, v, kc, vc, qm, stats, threshold, size;
+        int tokens, blocks;
+
+        explicit SolWorkspace(const ggml_tensor* dst) {
+            const auto* input  = dst->src[0];
+            const size_t heads = input->ne[2] * input->ne[3];
+            tokens             = GGML_PAD(input->ne[1], SOL_BLOCK);
+            blocks             = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
+            size               = 0;
+            q                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            k                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            v                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            kc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            vc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            qm                 = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
+            stats              = reserve(heads * 2 * SOL_DIM * sizeof(float));
+            threshold          = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
+        }
+
+        size_t reserve(size_t bytes) {
+            const size_t offset = size;
+            size += GGML_PAD(bytes, 256);
+            return offset;
+        }
+    };
+
+    static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
+        const int index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (index >= padded * SOL_DIM) {
+            return;
+        }
+        const int token   = index / SOL_DIM;
+        const int channel = index % SOL_DIM;
+        const size_t src  = size_t(blockIdx.y) * length * SOL_DIM + index;
+        const size_t dst  = size_t(blockIdx.y) * padded * SOL_DIM + index;
+        qb[dst]           = __float2bfloat16(token < length ? q[src] : 0.f);
+        kb[dst]           = __float2bfloat16(token < length ? k[src] : 0.f);
+        vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
+            __float2bfloat16(token < length ? v[src] : 0.f);
+    }
+
+    static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
+        const int channel = threadIdx.x;
+        const int block   = blockIdx.x;
+        const int count   = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
+        float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
+        for (int i = 0; i < count; ++i) {
+            const int token    = block * SOL_BLOCK + i;
+            const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
+            sum_q += __bfloat162float(q[index]);
+            sum_k += __bfloat162float(k[index]);
+            sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
+        }
+        kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
+            __float2bfloat16(count ? sum_k / count : 0.f);
+        vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
+        if (count) {
+            qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
+        }
+    }
+
+    static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
+        const int channel = threadIdx.x;
+        float sum = 0.f, sum_sq = 0.f;
+        for (int i = 0; i < blocks; ++i) {
+            const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
+            sum += value;
+            sum_sq += value * value;
+        }
+        const float mean                                            = sum / blocks;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel]           = mean;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
+    }
+
+    static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
+        const int channel = threadIdx.x;
+        const float q     = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
+        float mean        = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
+        float var         = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
+        mean              = warp_reduce_sum(mean);
+        var               = warp_reduce_sum(var);
+        __shared__ float means[4], vars[4];
+        if (channel % 32 == 0) {
+            means[channel / 32] = mean;
+            vars[channel / 32]  = var;
+        }
+        __syncthreads();
+        if (channel == 0) {
+            mean = means[0] + means[1] + means[2] + means[3];
+            var  = vars[0] + vars[1] + vars[2] + vars[3];
+            threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
+                mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
+        }
+    }
+
+    struct SolQTile {
+        uint32_t x[4];
+    };
+    struct SolKTile {
+        uint32_t x[2];
+    };
+    struct SolTile {
+        float x[4] = {};
+        static __device__ __forceinline__ int get_i(int i) {
+            return (i / 2) * 8 + threadIdx.x / 4;
+        }
+        static __device__ __forceinline__ int get_j(int i) {
+            return (threadIdx.x % 4) * 2 + i % 2;
+        }
+    };
+
+    static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+            tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 2; ++i)
+            tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
+#if __CUDA_ARCH__ >= 800
+        auto* dst = reinterpret_cast<uint32_t*>(output.x);
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
+            : "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
+#else
+        asm("trap;");
+#endif
+    }
+
+    static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+            scores[n] = SolTile{};
+#pragma unroll
+            for (int d = 0; d < 8; ++d) {
+                SolKTile kt;
+                sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
+                sol_mma(scores[n], q[d], kt);
+            }
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                scores[n].x[i] *= scale;
+            }
+        }
+    }
+
+    template <bool PROXY>
+    static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
+        float maximum[2] = {-INFINITY, -INFINITY};
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
+                                         : (start + column < length);
+                scores[n].x[i]   = valid ? scores[n].x[i] : -INFINITY;
+                maximum[i / 2]   = fmaxf(maximum[i / 2], scores[n].x[i]);
+            }
+        }
+        float alpha[2], sums[2] = {0.f, 0.f};
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
+            const float next = fmaxf(row_max[row], maximum[row]);
+            alpha[row]       = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
+            row_max[row]     = next;
+        }
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const float p    = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
+                const int count  = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
+                sums[i / 2] += p * count;
+                probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
+            }
+        }
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
+            row_sum[row] = row_sum[row] * alpha[row] + sums[row];
+        }
+        __syncwarp();
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                output[n].x[i] *= alpha[i / 2];
+            }
+#pragma unroll
+            for (int d = 0; d < 4; ++d) {
+                SolQTile pt;
+                SolKTile vt;
+                sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
+                sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
+                sol_mma(output[n], pt, vt);
+            }
+        }
+        __syncwarp();
+    }
+
+    static __global__ __launch_bounds__(128) void sol_forward(
+        const nv_bfloat16* q,
+        const nv_bfloat16* k,
+        const nv_bfloat16* v,
+        const nv_bfloat16* kc,
+        const nv_bfloat16* vc,
+        const float* thresholds,
+        float* dst,
+        int length,
+        int padded,
+        int padded_blocks,
+        int heads,
+        float scale) {
+#if __CUDA_ARCH__ >= 800
+        const int warp    = threadIdx.y;
+        const int lane    = threadIdx.x;
+        const int tid     = warp * 32 + lane;
+        const int blocks  = padded / SOL_BLOCK;
+        const int q_start = blockIdx.x * SOL_BLOCK;
+        const int q_len   = min(SOL_BLOCK, length - q_start);
+        const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
+        SolQTile qt[8];
+#pragma unroll
+        for (int d = 0; d < 8; ++d) {
+            sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
+        }
+        SolTile output[16];
+        float row_max[2] = {-INFINITY, -INFINITY};
+        float row_sum[2] = {0.f, 0.f};
+        __shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
+        __shared__ float routes[4][SOL_BLOCK];
+        __shared__ int exact[SOL_BLOCK];
+        const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
+        for (int group = 0; group < blocks; group += SOL_BLOCK) {
+            SolTile scores[8];
+            sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
+#pragma unroll
+            for (int n = 0; n < 8; ++n) {
+#pragma unroll
+                for (int column = 0; column < 2; ++column) {
+                    float sum = scores[n].x[column] + scores[n].x[column + 2];
+                    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+                    if (lane < 4) {
+                        routes[warp][n * 8 + 2 * lane + column] = sum;
+                    }
+                }
+            }
+            __syncthreads();
+            if (tid < SOL_BLOCK) {
+                const int block   = group + tid;
+                const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
+                exact[tid]        = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
+            }
+            __syncthreads();
+            sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                             vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
+                             padded_blocks, group, length, exact);
+            for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
+                if (!exact[block]) {
+                    continue;
+                }
+                const int start = (group + block) * SOL_BLOCK;
+                sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
+                sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                                  v + base + start, padded, start, length, nullptr);
+            }
+            __syncthreads();
+        }
+        const int head  = blockIdx.y % heads;
+        const int batch = blockIdx.y / heads;
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int token = q_start + warp * 16 + SolTile::get_i(i);
+                if (token < length) {
+                    const int channel = n * 8 + SolTile::get_j(i);
+                    dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
+                        __bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
+                }
+            }
+        }
+#else
+        asm("trap;");
+#endif
+    }
+
+    static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(op, &params))
+            return false;
+        const auto* q = op->src[0];
+        if (q->ne[2] * q->ne[3] > 65535 || GGML_PAD(q->ne[1], SOL_BLOCK) > INT_MAX / SOL_DIM)
+            return false;
+        for (const auto* tensor : {op->src[0], op->src[1], op->src[2]})
+            if (reinterpret_cast<uintptr_t>(tensor->data) & 3)
+                return false;
+        return !(reinterpret_cast<uintptr_t>(op->data) & 3);
+    }
+
+    static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+        *workspace = SolWorkspace(nodes[0]).size;
+        return GGML_STATUS_SUCCESS;
+    }
+    static void sol_destroy(void*, void*) {}
+
+    static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        auto* dst = nodes[0];
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(dst, &params))
+            return GGML_STATUS_FAILED;
+        const auto* q = dst->src[0];
+        const SolWorkspace work(dst);
+        char* data       = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        auto* qb         = reinterpret_cast<nv_bfloat16*>(data + work.q);
+        auto* kb         = reinterpret_cast<nv_bfloat16*>(data + work.k);
+        auto* vb         = reinterpret_cast<nv_bfloat16*>(data + work.v);
+        auto* kc         = reinterpret_cast<nv_bfloat16*>(data + work.kc);
+        auto* vc         = reinterpret_cast<nv_bfloat16*>(data + work.vc);
+        auto* qm         = reinterpret_cast<float*>(data + work.qm);
+        auto* stats      = reinterpret_cast<float*>(data + work.stats);
+        auto* threshold  = reinterpret_cast<float*>(data + work.threshold);
+        const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
+        const int blocks  = work.tokens / SOL_BLOCK;
+        const float scale = params.scale * 1.4426950408889634f;
+        const auto stream = static_cast<cudaStream_t>(context.stream);
+        sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
+            qb, kb, vb, length, work.tokens);
+        sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
+        sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
+        sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
+        sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
+                                                                           static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_sol_attention(ggml_backend_dev_t device) {
+        const auto info = cuda_device_info(device);
+        if (info.device < 0)
+            return;
+        int previous;
+        cudaDeviceProp properties;
+        cudaFuncAttributes attributes;
+        if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
+            cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
+            return;
+        const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
+        if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
+            std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
+            return;
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.sol_attn.cuda.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_sol_attn_op();
+        desc.priority             = 100;
+        desc.supports             = sol_supports;
+        desc.prepare              = sol_prepare;
+        desc.execute              = sol_execute;
+        desc.destroy              = sol_destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}
+#else
+namespace sd_cuda {
+    void register_sol_attention(ggml_backend_dev_t) {}
+}
+#endif

src/core/ggml_compat.cpp

@@ -0,0 +1,185 @@
+#include "core/ggml_compat.h"
+#include "core/ggml_extend_backend.h"
+
+#ifdef SD_GGML_GRAPH_COMPAT
+// Upstream ggml still requires its private graph layout for these operations.
+#include "ggml-impl.h"
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+int ggml_graph_n_leafs(const ggml_cgraph* graph) {
+    return graph->n_leafs;
+}
+
+ggml_tensor* ggml_graph_leaf(const ggml_cgraph* graph, int index) {
+    if (index < 0) {
+        index += graph->n_leafs;
+    }
+    GGML_ASSERT(index >= 0 && index < graph->n_leafs);
+    return graph->leafs[index];
+}
+
+void ggml_graph_add_leaf(ggml_cgraph* graph, ggml_tensor* tensor) {
+    GGML_ASSERT(tensor && graph->n_leafs < graph->size);
+    graph->leafs[graph->n_leafs++] = tensor;
+}
+
+int32_t ggml_graph_use_count(const ggml_cgraph* graph, const ggml_tensor* tensor) {
+    if (!graph || !tensor || !graph->visited_hash_set.size) {
+        return 0;
+    }
+    const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
+    return index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
+               ? graph->use_counts[index]
+               : 0;
+}
+
+bool ggml_graph_cpy_mapped(const ggml_cgraph* src, ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data) {
+    if (!src || !dst || src == dst || !map || dst->grads || dst->grad_accs ||
+        src->n_nodes > dst->size || src->n_leafs > dst->size) {
+        return false;
+    }
+    std::unordered_map<const ggml_tensor*, ggml_tensor*> mapping;
+    std::unordered_set<ggml_tensor*> mapped;
+    std::vector<const ggml_tensor*> sources;
+    bool valid = true;
+    auto remap = [&](const ggml_tensor* source) {
+        const auto found = mapping.find(source);
+        if (found != mapping.end()) {
+            return found->second;
+        }
+        auto* tensor = map(source, user_data);
+        if (tensor && !mapped.insert(tensor).second) {
+            valid = false;
+        }
+        mapping.emplace(source, tensor);
+        sources.push_back(source);
+        return tensor;
+    };
+    std::vector<ggml_tensor*> nodes(src->n_nodes), leafs(src->n_leafs);
+    for (int i = 0; i < src->n_nodes; ++i) {
+        if (!(nodes[i] = remap(src->nodes[i])) || !valid) {
+            return false;
+        }
+    }
+    for (int i = 0; i < src->n_leafs; ++i) {
+        if (!(leafs[i] = remap(src->leafs[i])) || !valid) {
+            return false;
+        }
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (size_t i = 0; i < sources.size(); ++i) {
+        const auto* source = sources[i];
+        if (source->view_src) {
+            remap(source->view_src);
+        }
+        for (const auto* input : source->src) {
+            if (input) {
+                remap(input);
+            }
+        }
+        if (!valid) {
+            return false;
+        }
+        auto* tensor = mapping.at(source);
+        if (!tensor || !src->visited_hash_set.size) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&src->visited_hash_set, source);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, index)) {
+            continue;
+        }
+        if (counts.size() >= dst->visited_hash_set.size) {
+            return false;
+        }
+        counts.emplace_back(tensor, src->use_counts[index]);
+    }
+    ggml_graph_clear(dst);
+    dst->order = src->order;
+    dst->uid   = 0;
+    for (auto* tensor : nodes) {
+        ggml_graph_add_node(dst, tensor);
+    }
+    for (auto* tensor : leafs) {
+        ggml_graph_add_leaf(dst, tensor);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+bool ggml_graph_cpy_use_counts(const ggml_cgraph* src, ggml_cgraph* dst) {
+    if (!src || !dst || !src->visited_hash_set.size || !dst->visited_hash_set.size) {
+        return false;
+    }
+    std::unordered_set<ggml_tensor*> represented;
+    for (int i = 0; i < dst->n_nodes; ++i) {
+        represented.insert(dst->nodes[i]);
+    }
+    for (int i = 0; i < dst->n_leafs; ++i) {
+        represented.insert(dst->leafs[i]);
+    }
+    size_t occupied = 0;
+    for (size_t i = 0; i < dst->visited_hash_set.size; ++i) {
+        occupied += ggml_bitset_get(dst->visited_hash_set.used, i);
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (auto* tensor : represented) {
+        if (!tensor) {
+            return false;
+        }
+        const size_t source_index = ggml_hash_find(&src->visited_hash_set, tensor);
+        if (source_index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, source_index)) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&dst->visited_hash_set, tensor);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(dst->visited_hash_set.used, index)) {
+            if (++occupied > dst->visited_hash_set.size) {
+                return false;
+            }
+        }
+        counts.emplace_back(tensor, src->use_counts[source_index]);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_find_or_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend, ggml_cgraph* graph, int first, int last) {
+    if (!backend || !graph || first < 0 || last < first || last > graph->n_nodes) {
+        return GGML_STATUS_FAILED;
+    }
+    if (first == last) {
+        return GGML_STATUS_SUCCESS;
+    }
+    auto view = ggml_graph_view(graph, first, last);
+    return ggml_backend_graph_compute_async(backend, &view);
+}
+#endif
+
+#ifdef SD_USE_UPSTREAM_GGML
+#include "ggml-backend-impl.h"
+#endif
+
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void* ptr, size_t size, size_t max_tensor_size) {
+    auto buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
+#ifdef SD_USE_UPSTREAM_GGML
+    if (buffer && !buffer->context) {
+        auto reg = ggml_backend_dev_backend_reg(device);
+        if (reg && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
+            // Older Metal backends wrap failed mappings and cannot free their null context.
+            buffer->iface.free_buffer = nullptr;
+            ggml_backend_buffer_free(buffer);
+            return nullptr;
+        }
+    }
+#endif
+    return buffer;
+}

src/core/ggml_compat.h

@@ -0,0 +1,43 @@
+#ifndef __SD_CORE_GGML_COMPAT_H__
+#define __SD_CORE_GGML_COMPAT_H__
+
+#include "ggml-backend.h"
+
+#if __has_include("ggml-graph.h")
+#include "ggml-graph.h"
+#elif defined(SD_USE_UPSTREAM_GGML)
+#define SD_GGML_GRAPH_COMPAT
+
+int ggml_graph_n_leafs(const struct ggml_cgraph* graph);
+struct ggml_tensor* ggml_graph_leaf(const struct ggml_cgraph* graph, int index);
+// Like ggml_graph_add_node, appends metadata without traversing tensor dependencies.
+void ggml_graph_add_leaf(struct ggml_cgraph* graph, struct ggml_tensor* tensor);
+// Returns zero when the tensor has no recorded use count. Includes uses outside graph views.
+int32_t ggml_graph_use_count(const struct ggml_cgraph* graph, const struct ggml_tensor* tensor);
+
+typedef struct ggml_tensor* (*ggml_graph_tensor_map_fn)(const struct ggml_tensor* tensor, void* user_data);
+
+// Copies forward node/leaf order and recorded use counts, including external consumers.
+// The callback must be injective and may return NULL for dependencies outside the
+// node/leaf lists. It is called once per tensor in their dependency closure.
+// Tensor storage and edges are mapped by the caller.
+// The destination must have no gradients and enough capacity. Gradients are not copied.
+// Invalid mappings leave the destination unchanged. src and dst must be distinct.
+bool ggml_graph_cpy_mapped(const struct ggml_cgraph* src, struct ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data);
+
+// Copies recorded counts for tensors in dst's node/leaf lists. Other metadata is unchanged.
+// This preserves external consumers when constructing an inference graph segment.
+bool ggml_graph_cpy_use_counts(const struct ggml_cgraph* src, struct ggml_cgraph* dst);
+
+// Executes nodes [first, last) with the original graph's use counts. No graph ownership
+// is transferred. Tensor storage must already be allocated; invalid bounds fail.
+enum ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend,
+                                                        struct ggml_cgraph* graph,
+                                                        int first,
+                                                        int last);
+
+#else
+#error "Patched ggml requires ggml-graph.h; select SD_USE_UPSTREAM_GGML for upstream compatibility"
+#endif
+
+#endif

src/core/ggml_extend_backend.cpp

@@ -18,9 +18,8 @@
 #include <hip/hip_runtime_api.h>
 #endif
 
+#include "core/ggml_compat.h"
 #include "core/util.h"
-#include "ggml-backend-impl.h"
-#include "ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -451,24 +450,6 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
     return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
 }
 
-ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
-                                                          void* ptr,
-                                                          size_t size,
-                                                          size_t max_tensor_size) {
-    ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
-    if (buffer != nullptr && buffer->context == nullptr) {
-        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
-        if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
-            // Metal can wrap a failed mapping in a non-null buffer. Its free callback also
-            // dereferences the missing context, so only release the outer buffer.
-            buffer->iface.free_buffer = nullptr;
-            ggml_backend_buffer_free(buffer);
-            return nullptr;
-        }
-    }
-    return buffer;
-}
-
 static int cuda_compiled_arch(ggml_backend_t backend) {
 #ifdef SD_USE_CUDA
     if (!sd_backend_is(backend, "CUDA")) {
@@ -644,23 +625,6 @@ void sd_backend_clear_buffer_cache(ggml_backend_t backend) {
     }
 }
 
-static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
-    ggml_cgraph cgraph = {
-        /*.size             =*/0,
-        /*.n_nodes          =*/i1 - i0,
-        /*.n_leafs          =*/0,
-        /*.nodes            =*/cgraph0->nodes + i0,
-        /*.grads            =*/nullptr,
-        /*.grad_accs        =*/nullptr,
-        /*.leafs            =*/nullptr,
-        /*.use_counts       =*/cgraph0->use_counts,
-        /*.visited_hash_set =*/cgraph0->visited_hash_set,
-        /*.order            =*/cgraph0->order,
-        /*.uid              =*/0,
-    };
-    return cgraph;
-}
-
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         ggml_cgraph* gf,
                                                         sd_graph_eval_callback_t callback_eval,
@@ -683,8 +647,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
             need = callback_eval(t, true, callback_eval_user_data);
         }
 
-        ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
-        status         = ggml_backend_graph_compute_async(backend, &gv);
+        status = ggml_backend_graph_compute_async_range(backend, gf, j0, j1 + 1);
         if (status != GGML_STATUS_SUCCESS) {
             break;
         }

src/core/ggml_extensions.cpp

@@ -277,6 +277,58 @@ ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, g
     return ggml_ext_op_create(ctx, sd_ggml_sage_attn_op(), inputs, 3, &params);
 }
 
+static bool sol_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+    if (count != 3 || !attributes)
+        return false;
+    SDSolAttentionParams params;
+    std::memcpy(&params, attributes, sizeof(params));
+    if (!std::isfinite(params.scale) || params.scale <= 0.f || !std::isfinite(params.tau))
+        return false;
+    const auto* q = inputs[0];
+    if (!q || q->ne[0] != 128)
+        return false;
+    for (int i = 0; i < count; ++i) {
+        const auto* t = inputs[i];
+        if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || !ggml_are_same_shape(q, t))
+            return false;
+    }
+    int64_t elements = 1;
+    for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+        if (q->ne[d] <= 0 || q->ne[d] > INT64_MAX / int64_t(sizeof(float)) / elements)
+            return false;
+        elements *= q->ne[d];
+    }
+    *out = {GGML_TYPE_F32, {q->ne[0], q->ne[2], q->ne[1], q->ne[3]}};
+    return true;
+}
+
+ggml_ext_op_t sd_ggml_sol_attn_op() {
+    static const auto op = [] {
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.sol_attn.v1", sizeof(SDSolAttentionParams), sol_attn_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params) {
+    if (!node || ggml_ext_op_get(node) != sd_ggml_sol_attn_op() || node->src[3] || !ggml_is_contiguous(node))
+        return false;
+    ggml_ext_tensor_desc output;
+    const auto* attributes = ggml_ext_op_attributes(node);
+    if (!sol_attn_infer(node->src, 3, attributes, &output) || node->type != output.type || std::memcmp(node->ne, output.ne, sizeof(output.ne)))
+        return false;
+    std::memcpy(params, attributes, sizeof(*params));
+    return true;
+}
+
+ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau) {
+    ggml_tensor* inputs[]             = {q, k, v};
+    const SDSolAttentionParams params = {scale, tau};
+    return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, &params);
+}
+
 void sd_ggml_register_extensions(ggml_backend_t backend) {
     if (!backend) {
         return;

src/core/ggml_extensions.h

@@ -51,6 +51,16 @@ bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams*
 // Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout.
 // Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference.
 ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode);
+
+struct SDSolAttentionParams {
+    float scale;
+    float tau;
+};
+
+ggml_ext_op_t sd_ggml_sol_attn_op();
+bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params);
+// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
+ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau);
 #endif
 
 #endif

src/core/ggml_graph_cut.cpp

@@ -16,7 +16,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 
-#include "ggml-impl.h"
+#include "core/ggml_compat.h"
 
 namespace sd::ggml_graph_cut {
 
@@ -35,8 +35,8 @@ namespace sd::ggml_graph_cut {
     static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(tensor != nullptr);
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            if (gf->leafs[i] == tensor) {
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            if (ggml_graph_leaf(gf, i) == tensor) {
                 return i;
             }
         }
@@ -389,7 +389,7 @@ namespace sd::ggml_graph_cut {
                         return fail(segment_label + " references an unavailable cut node " +
                                     std::to_string(input.node_index));
                     }
-                } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) {
+                } else if (input.leaf_index < 0 || input.leaf_index >= ggml_graph_n_leafs(gf)) {
                     return fail(segment_label + " references an invalid leaf " +
                                 std::to_string(input.leaf_index));
                 }
@@ -444,15 +444,15 @@ namespace sd::ggml_graph_cut {
 
     int leaf_count(ggml_cgraph* gf) {
         GGML_ASSERT(gf != nullptr);
-        return gf->n_leafs;
+        return ggml_graph_n_leafs(gf);
     }
 
     ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
         GGML_ASSERT(gf != nullptr);
-        if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
+        if (leaf_index < 0 || leaf_index >= ggml_graph_n_leafs(gf)) {
             return nullptr;
         }
-        return gf->leafs[leaf_index];
+        return ggml_graph_leaf(gf, leaf_index);
     }
 
     ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
@@ -548,14 +548,14 @@ namespace sd::ggml_graph_cut {
     };
 
     std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
-        const size_t graph_size = static_cast<size_t>(graph->n_leafs) + graph->n_nodes;
+        const size_t graph_size = static_cast<size_t>(ggml_graph_n_leafs(graph)) + ggml_graph_n_nodes(graph);
         GraphLayoutTensors layout_tensors(graph_size);
         const auto& tensors = layout_tensors.tensors;
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            layout_tensors.add(graph->leafs[i]);
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            layout_tensors.add(ggml_graph_leaf(graph, i));
         }
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            layout_tensors.add(graph->nodes[i]);
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            layout_tensors.add(ggml_graph_node(graph, i));
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
             layout_tensors.add(tensors[i]->view_src);
@@ -567,13 +567,13 @@ namespace sd::ggml_graph_cut {
         const size_t tensor_fields = 5 + 2 * GGML_MAX_DIMS + GGML_MAX_SRC +
                                      GGML_MAX_OP_PARAMS / sizeof(int32_t) + (include_bindings ? 2 : 0);
         signature.reserve(2 + graph_size + tensors.size() * tensor_fields);
-        signature.push_back(graph->n_nodes);
-        signature.push_back(graph->n_leafs);
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            signature.push_back(layout_tensors.index(graph->leafs[i]));
+        signature.push_back(ggml_graph_n_nodes(graph));
+        signature.push_back(ggml_graph_n_leafs(graph));
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            signature.push_back(layout_tensors.index(ggml_graph_leaf(graph, i)));
         }
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            signature.push_back(layout_tensors.index(graph->nodes[i]));
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            signature.push_back(layout_tensors.index(ggml_graph_node(graph, i)));
         }
         for (auto tensor : tensors) {
             signature.push_back(tensor->op);
@@ -606,18 +606,18 @@ namespace sd::ggml_graph_cut {
                                    const Plan& plan,
                                    const std::vector<uint64_t>& layout) {
         GGML_ASSERT(gf != nullptr);
-        if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
+        if (plan.leaf_names.size() != static_cast<size_t>(ggml_graph_n_leafs(gf)) ||
             plan.layout != layout) {
             return false;
         }
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            if (plan.leaf_names[i] != gf->leafs[i]->name) {
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            if (plan.leaf_names[i] != ggml_graph_leaf(gf, i)->name) {
                 return false;
             }
         }
         size_t cut_index = 0;
-        for (int i = 0; i < gf->n_nodes; ++i) {
-            auto node = gf->nodes[i];
+        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
+            auto node = ggml_graph_node(gf, i);
             if (is_graph_cut_tensor(node)) {
                 if (cut_index >= plan.cut_markers.size() ||
                     plan.cut_markers[cut_index].first != i || plan.cut_markers[cut_index].second != node->name) {
@@ -654,7 +654,7 @@ namespace sd::ggml_graph_cut {
             }
             return ggml_graph_node(gf, input_ref.node_index);
         }
-        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
+        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= ggml_graph_n_leafs(gf)) {
             return nullptr;
         }
         return leaf_tensor(gf, input_ref.leaf_index);
@@ -735,8 +735,7 @@ namespace sd::ggml_graph_cut {
         GGML_ASSERT(segment_graph != nullptr);
 
         for (ggml_tensor* leaf : leaves) {
-            GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
-            segment_graph->leafs[segment_graph->n_leafs++] = leaf;
+            ggml_graph_add_leaf(segment_graph, leaf);
         }
 
         for (int output_node_index : segment.output_node_indices) {
@@ -753,14 +752,7 @@ namespace sd::ggml_graph_cut {
         for (int node_idx : segment.internal_node_indices) {
             ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
         }
-        // Fusion checks need the original use counts, including consumers outside this segment.
-        for (ggml_tensor* tensor : represented) {
-            const size_t source_index = ggml_hash_find(&gf->visited_hash_set, tensor);
-            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(gf->visited_hash_set.used, source_index)) {
-                const size_t index               = ggml_hash_insert(&segment_graph->visited_hash_set, tensor);
-                segment_graph->use_counts[index] = gf->use_counts[source_index];
-            }
-        }
+        GGML_ASSERT(ggml_graph_cpy_use_counts(gf, segment_graph));
         *graph_ctx_out = graph_ctx;
         return segment_graph;
     }
@@ -896,8 +888,8 @@ namespace sd::ggml_graph_cut {
             return plan;
         }
         plan.layout = graph_layout(gf, false);
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            plan.leaf_names.emplace_back(gf->leafs[i]->name);
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            plan.leaf_names.emplace_back(ggml_graph_leaf(gf, i)->name);
         }
         plan.compute_buffer_size =
             measure_graph_compute_buffer(backend, gf, params_tensor_set);