From 8d3587965a57c2acd54fd3291de503afad9defeb Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 1 Nov 2026 03:03:53 +0800 Subject: [PATCH] refactor: externalize Sol-Attn and use public graph APIs --- cmake/ggml.cmake | 22 +- docs/ggml_extensions.md | 38 ++- ggml | 2 +- src/core/compute_workspace.cpp | 47 ++-- src/core/cuda/CMakeLists.txt | 4 +- src/core/cuda/registration.cu | 2 + src/core/cuda/sol_attention.cu | 409 +++++++++++++++++++++++++++++++ src/core/ggml_compat.cpp | 185 ++++++++++++++ src/core/ggml_compat.h | 43 ++++ src/core/ggml_extend_backend.cpp | 41 +--- src/core/ggml_extensions.cpp | 52 ++++ src/core/ggml_extensions.h | 10 + src/core/ggml_graph_cut.cpp | 64 +++-- 13 files changed, 798 insertions(+), 121 deletions(-) create mode 100644 src/core/cuda/sol_attention.cu create mode 100644 src/core/ggml_compat.cpp create mode 100644 src/core/ggml_compat.h diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index 39d77eae..55d0067b 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -9,17 +9,19 @@ if(NOT TARGET ggml) add_library(ggml ALIAS ggml::ggml) endif() -get_target_property(sd_ggml_imported ggml IMPORTED) -if(sd_ggml_imported) - set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src") -else() - get_target_property(sd_ggml_private_include ggml SOURCE_DIR) +if(SD_USE_UPSTREAM_GGML) + get_target_property(sd_ggml_imported ggml IMPORTED) + if(sd_ggml_imported) + set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src") + else() + get_target_property(sd_ggml_private_include ggml SOURCE_DIR) + endif() + if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h") + message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the source tree matching the selected ggml library (ggml-impl.h is required).") + endif() + target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}") + set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}") endif() -if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h") - message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the source tree matching the selected ggml library (ggml-impl.h is required).") -endif() -target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}") -set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}") if(SD_USE_UPSTREAM_GGML) target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index c196f0ef..6787c443 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -399,8 +399,8 @@ The corresponding Vulkan shader sources and native fusion rules have been removed from ggml. Unsupported custom nodes retain scheduler fallback to another supported backend. -Remaining model-specific CUDA paths still need auditing or migration. This -implementation does not add dynamic dtypes. +RoPE Apply, INT8/ConvRot, CuTe attention, SageAttention, and Sol-Attn use +application-owned implementations. This interface does not add dynamic dtypes. `src/core/vulkan/int8.cpp` owns Vulkan INT8 GEMM, ConvRot, and their fusion rules. The Vulkan compute metadata service exposes shared-memory limits, subgroup @@ -429,3 +429,37 @@ the application's compiled CUDA target as well as the physical device. CPU, HIP, MUSA, and Vulkan do not implement this NVIDIA-specific kernel. Existing model-level capability checks retain flash/default attention for unsupported layers; a custom Sage node itself does not silently change algorithms. + +## Sol-Attn + +The existing `ggml_sol_attn` entry point now lives in the common sd.cpp extension +files and constructs `sd.sol_attn.v1`. It accepts contiguous F32 Q/K/V tensors +with shape `[128, tokens, heads, batch]`, a positive finite scale, and finite tau. +The output remains F32 `[128, heads, tokens, batch]`. No current sd.cpp model or +CLI path calls this entry point; the migration preserves its low-level behavior. + +`src/core/cuda/sol_attention.cu` owns packing, block summaries, threshold +selection, and the BF16 tensor-core kernel. Temporary storage is separate +scheduler workspace, so the output allocation contains only the output tensor. +The kernel uses the public backend stream service and direct CUDA/PTX operations. +Registration requires CUDA 12 or newer, an SM80-or-newer device, and an application +kernel compiled for SM80 or newer. CPU, HIP, MUSA, and Vulkan do not advertise +support. ggml no longer contains a Sol-Attn opcode, constructor, allocation rule, +or backend kernel. + +## Public graph access + +`ggml-graph.h` exposes leaf enumeration, append-only leaf construction, recorded +use counts, mapped inference-graph copying, use-count copying for segments, and +backend execution of a node range. Mapped copies retain node order and consumer +counts from the original graph, including consumers outside a graph view. The +caller owns tensor copies and remaps their storage and edges; invalid mappings +leave the destination graph unchanged. Range execution uses already allocated +tensors and retains the original graph's consumer counts. + +Workspace measurement, graph cuts, and evaluation callbacks use these APIs. +Normal sd.cpp builds require only ggml's public headers. Compatibility with an +upstream ggml lacking this header is isolated in `src/core/ggml_compat.cpp`; that +optional build still requires matching upstream private headers. The compatibility +layer also retains the workaround for old Metal backends wrapping failed host +mappings. The patched Metal backend returns null directly when mapping fails. diff --git a/ggml b/ggml index 62da2f1b..50ebf58e 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 62da2f1b7e2ff90241dc48fb0fa3e9fd003706f5 +Subproject commit 50ebf58e14e0974ba043183c8c2b2c8bf6657a2c diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index 1c11c12d..eaf13184 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -6,9 +6,9 @@ #include #include +#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(graph->n_nodes + graph->n_leafs) + 8; + const size_t capacity = static_cast(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(&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(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(©_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*>(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) { diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 9a610a49..e6834c66 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/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) diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index f13d94f2..d03b114c 100644 --- a/src/core/cuda/registration.cu +++ b/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); diff --git a/src/core/cuda/sol_attention.cu b/src/core/cuda/sol_attention.cu new file mode 100644 index 00000000..0b580abd --- /dev/null +++ b/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 +#include +#include + +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(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(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(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(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 + 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(probabilities + d * 16), SOL_BLOCK / 2); + sol_load(vt, reinterpret_cast(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(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(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(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, ¶ms)) + 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(tensor->data) & 3) + return false; + return !(reinterpret_cast(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, ¶ms)) + return GGML_STATUS_FAILED; + const auto* q = dst->src[0]; + const SolWorkspace work(dst); + char* data = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + auto* qb = reinterpret_cast(data + work.q); + auto* kb = reinterpret_cast(data + work.k); + auto* vb = reinterpret_cast(data + work.v); + auto* kc = reinterpret_cast(data + work.kc); + auto* vc = reinterpret_cast(data + work.vc); + auto* qm = reinterpret_cast(data + work.qm); + auto* stats = reinterpret_cast(data + work.stats); + auto* threshold = reinterpret_cast(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(context.stream); + sol_pack<<>>( + static_cast(q->data), static_cast(dst->src[1]->data), static_cast(dst->src[2]->data), + qb, kb, vb, length, work.tokens); + sol_summaries<<>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks); + sol_statistics<<>>(kc, stats, blocks, work.blocks); + sol_thresholds<<>>(qm, stats, threshold, blocks, scale, params.tau); + sol_forward<<>>(qb, kb, vb, kc, vc, threshold, + static_cast(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 diff --git a/src/core/ggml_compat.cpp b/src/core/ggml_compat.cpp new file mode 100644 index 00000000..061b3370 --- /dev/null +++ b/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 +#include +#include +#include + +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 mapping; + std::unordered_set mapped; + std::vector 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 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> 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 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> 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; +} diff --git a/src/core/ggml_compat.h b/src/core/ggml_compat.h new file mode 100644 index 00000000..9f64262f --- /dev/null +++ b/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 diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 1ea86a8c..47fe80ff 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -18,9 +18,8 @@ #include #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; } diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index e0214fe4..1f0aa1de 100644 --- a/src/core/ggml_extensions.cpp +++ b/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, ¶ms); } +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(¶ms, 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, ¶ms); +} + void sd_ggml_register_extensions(ggml_backend_t backend) { if (!backend) { return; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index d069f10e..49b0e760 100644 --- a/src/core/ggml_extensions.h +++ b/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 diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp index 8454d435..fc6d8c63 100644 --- a/src/core/ggml_graph_cut.cpp +++ b/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 graph_layout(ggml_cgraph* graph, bool include_bindings) { - const size_t graph_size = static_cast(graph->n_leafs) + graph->n_nodes; + const size_t graph_size = static_cast(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& layout) { GGML_ASSERT(gf != nullptr); - if (plan.leaf_names.size() != static_cast(gf->n_leafs) || + if (plan.leaf_names.size() != static_cast(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);