Commit 5f8c19f for stable-diffusion.cpp
commit 5f8c19f98ddb84e301bb8725b5260ec79e768e22
Author: leejet <leejet714@gmail.com>
Date: Sun Nov 8 12:28:44 2026 +0800
perf: reduce fused compute buffers (#2290)
diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md
index 7707598..580ef58 100644
--- a/docs/ggml_extensions.md
+++ b/docs/ggml_extensions.md
@@ -89,8 +89,13 @@ again after reordering.
External inputs stay allocated through the final node. Native graph optimization
and execution operate on the remaining contiguous ranges, so native fusion cannot
-consume nodes selected by an extension. Intermediate allocations are retained
-in this initial implementation; only their computation is removed.
+consume nodes selected by an extension. Fused intermediates are excluded from
+allocation through the backend's `skip_alloc` callback and
+`ggml_gallocr_set_skipped_tensors`. Their metadata and graph edges remain available,
+but execution must not access their data or buffers. Storage backing the final
+output is retained, including an internal owner used by an in-place final write.
+The skipped set is replaced for each allocation graph; changing which tensors are skipped
+invalidates cached allocation assignments even when tensor shapes are unchanged.
Automatic extension fusion is enabled in scheduler execution with one pipeline
copy and without evaluation callbacks. Multiple pipeline copies and evaluation
@@ -137,6 +142,17 @@ Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; caller
must use the representation appropriate to their backend. Vulkan extensions use
the command recording services described below.
+CUDA INT8 Linear, RoPE, ConvRot, RMS affine and INT8 gated residual kernels omit output
+scratch when allocation guarantees that a fresh owning output cannot overlap
+live inputs. Existing outputs and views retain overlap protection. Internal
+GEMM and quantization workspace is still reserved when required; kernels also
+handle a null workspace when the entire step needs no scratch.
+
+Compute workspace measurement preserves the storage identities and offsets of
+external tensors, including views. Separate allocations must not share a synthetic
+buffer identity: that can invent overlap, overestimate scratch, and repeatedly
+rebuild a workspace whose actual allocation is smaller than the estimate.
+
## sd.cpp integration
`src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas,
@@ -202,6 +218,10 @@ ConvRot output fusion. These kernels use only the public extension API, borrowed
CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or
private conversion helpers. The existing CUDA fusion and CuTe disable environment
variables are honored.
+Aligned contiguous FP16 K/V inputs are read directly without reserving conversion
+buffers. Fresh F32 outputs are written by the attention epilogue, retaining the
+original FP16 rounding before output scaling. Packed ConvRot and preallocated or
+view outputs retain the intermediate half buffer.
`src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues:
`GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy,
@@ -262,12 +282,21 @@ small-row policy, Stream-K scheduling, Ada column bands, and separate scaling
order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row
concatenation for custom GEMM nodes. Weight scales remain explicit
inputs, including the gate/up scale permutation recorded on every SwiGLU launch.
+Two independent INT8 GEMMs sharing a quantized input can also fuse with split
+SwiGLU. A paired-column iterator reads the separate gate/up weight allocations,
+and the GEMM epilogue writes only the activated output. Both intermediate
+activations are omitted through the generic fusion allocation interface; weights
+are neither copied nor concatenated. This path requires bias-free branches,
+compatible layouts and alignments, and no external intermediate consumers.
+Only interleaved weight scales and any Stream-K workspace remain temporary.
Plans reserve bounded Stream-K workspace, cache occupancy on first execution,
and initialize barriers on every launch, so CUDA Graph replay observes updated
inputs and scales. External consumers retain ordinary execution; overlapping
-outputs use workspace. The CUTLASS SwiGLU/ConvRot chain can write its packed
-output directly after GEMM, because quantization reads only its workspace
-intermediate. Shapes or alignments unsuitable for CUTLASS use cuBLAS.
+outputs use workspace. For fresh outputs, the CUTLASS SwiGLU epilogue writes
+an allocator-managed tensor that the subsequent ConvRot launch consumes. This
+keeps the same GPU launches while allowing its storage to be reused by other
+operations. The combined SwiGLU/ConvRot path retains its workspace intermediate
+for other cases. Shapes or alignments unsuitable for CUTLASS use cuBLAS.
Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to
preserve their previous accumulation behavior. Fusion disable and force-cuBLAS
settings are honored. No CUTLASS headers are needed for the cuBLAS path.
diff --git a/ggml b/ggml
index 135ce1e..3fc478e 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820
+Subproject commit 3fc478e6191aadf51ff7cf7bef15afe6ebeb14b1
diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp
index 4e53175..cf9d04c 100644
--- a/src/core/compute_workspace.cpp
+++ b/src/core/compute_workspace.cpp
@@ -147,8 +147,8 @@ namespace sd {
append(type);
}
} else {
- const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
- append(alignment_offset);
+ append(tensor->buffer);
+ append(tensor->data);
}
const int32_t uses = ggml_graph_use_count(graph, tensor);
append(uses);
@@ -164,7 +164,7 @@ namespace sd {
return {};
}
std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
- std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
+ std::map<std::pair<ggml_backend_t, const void*>, ggml_backend_buffer_t> external_buffers;
for (auto tensor : tensors) {
auto copy = ggml_dup_tensor(context, tensor);
*copy = *tensor;
@@ -179,7 +179,9 @@ namespace sd {
}
auto external = external_backend(source);
if (external != nullptr && source->view_src == nullptr) {
- auto& buffer = external_buffers[external];
+ // Preserve distinct external allocations so overlap checks do not invent aliases.
+ const void* storage = source->buffer ? static_cast<const void*>(source->buffer) : source;
+ auto& buffer = external_buffers[{external, storage}];
if (buffer == nullptr) {
buffer = ggml_backend_alloc_buffer(external, 0);
GGML_ASSERT(buffer != nullptr);
@@ -191,6 +193,18 @@ namespace sd {
copy->extra = nullptr;
}
}
+ for (const auto& entry : copies) {
+ const auto* source = entry.first;
+ if (source->view_src && source->buffer) {
+ const auto* owner = source->view_src;
+ while (owner->view_src) {
+ owner = owner->view_src;
+ }
+ if (external_backend(owner) != nullptr) {
+ entry.second->buffer = copies.at(owner)->buffer;
+ }
+ }
+ }
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);
diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu
index cf78961..26920ae 100644
--- a/src/core/cuda/convrot.cu
+++ b/src/core/cuda/convrot.cu
@@ -204,7 +204,8 @@ namespace SD_CUDA_NAMESPACE {
plan->device = cuda_device_info(ggml_backend_get_device(backend));
if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
return GGML_STATUS_FAILED;
- *workspace = ggml_nbytes(nodes[count - 1]);
+ const auto* dst = nodes[count - 1];
+ *workspace = !dst->data && !dst->buffer && !dst->view_src ? 0 : ggml_nbytes(dst);
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu
index 1eaefa4..32bc891 100644
--- a/src/core/cuda/cute_attention.cu
+++ b/src/core/cuda/cute_attention.cu
@@ -237,6 +237,9 @@ namespace sd_fattn_cute {
const cutlass::half_t *q, *k, *v;
cutlass::half_t* o;
int nq, nk, heads, kv_heads, dim;
+ float* output_f32 = nullptr;
+ float output_scale = 1.0f, output_bias = 0.0f;
+ bool scaled_output = false;
};
template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
@@ -333,18 +336,34 @@ namespace sd_fattn_cute {
auto tOrO = threadO.retile_S(rO);
auto tOsO = threadO.partition_D(sO);
cute::copy(copyO, tOrO, tOsO);
- auto mO = make_tensor(make_gmem_ptr(p.o + int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim),
- make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
- auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
typename T::OutputCopy output_copy;
auto global_threadO = output_copy.get_thread_slice(tid);
auto from_sO = global_threadO.partition_S(sO);
- auto to_gO = global_threadO.partition_D(gO);
auto output_id = global_threadO.partition_D(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
__syncthreads();
- auto output_regs = make_tensor<Element>(shape(to_gO));
+ auto output_regs = make_tensor<Element>(shape(from_sO));
cute::copy(output_copy, from_sO, output_regs);
- copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+ const int64_t output_offset = int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim;
+ if (p.output_f32) {
+ auto float_copy = make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, float>{},
+ typename T::GlobalLayout{}, Layout<Shape<_1, _8>>{});
+ auto mO = make_tensor(make_gmem_ptr(p.output_f32 + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+ auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+ auto to_gO = float_copy.get_thread_slice(tid).partition_D(gO);
+ auto values = make_tensor<float>(shape(output_regs));
+#pragma unroll
+ for (int i = 0; i < size(values); ++i) {
+ // Preserve the half-rounded attention result before applying the F32 epilogue.
+ const float value = float(output_regs(i));
+ values(i) = p.scaled_output ? __fmaf_rn(value, p.output_scale, p.output_bias) : value;
+ }
+ copy_rows<false, false, FullDim>(float_copy, values, to_gO, output_id, p.nq - m_block * BM, p.dim);
+ } else {
+ auto mO = make_tensor(make_gmem_ptr(p.o + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+ auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+ auto to_gO = global_threadO.partition_D(gO);
+ copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+ }
}
}
@@ -625,17 +644,34 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
struct AttentionPlan {
int index;
size_t offsets[4];
+ bool direct_output;
};
+static bool attention_direct_input(const ggml_tensor* src) {
+ size_t offset = 0;
+ const auto* owner = src;
+ while (!owner->data && owner->view_src) {
+ offset += owner->view_offs;
+ owner = owner->view_src;
+ }
+ return src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && (uintptr_t(owner->data) + offset) % 16 == 0;
+}
+
static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
auto* plan = new AttentionPlan{};
*state = plan;
plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
: 0;
GGML_ASSERT(plan->index < count);
- const auto* attn = nodes[plan->index];
- size_t size = 0;
+ const auto* attn = nodes[plan->index];
+ const auto* dst = nodes[count - 1];
+ plan->direct_output = dst->type == GGML_TYPE_F32 && !dst->data && !dst->buffer && !dst->view_src;
+ size_t size = 0;
for (int i = 0; i < 4; ++i) {
+ if ((i > 0 && i < 3 && attention_direct_input(attn->src[i])) || (i == 3 && plan->direct_output)) {
+ plan->offsets[i] = SIZE_MAX;
+ continue;
+ }
plan->offsets[i] = size;
const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half);
size += (bytes + 255) / 256 * 256;
@@ -681,7 +717,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
auto* base = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
half* work[4];
for (int i = 0; i < 4; ++i)
- work[i] = reinterpret_cast<half*>(base + plan.offsets[i]);
+ work[i] = plan.offsets[i] == SIZE_MAX ? nullptr : reinterpret_cast<half*>(base + plan.offsets[i]);
const half* inputs[3] = {work[0], work[1], work[2]};
for (int i = 0; i < 3; ++i) {
const auto* src = attn->src[i];
@@ -695,6 +731,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
} else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
inputs[i] = static_cast<const half*>(src->data);
} else {
+ if (!work[i])
+ return GGML_STATUS_FAILED;
auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
@@ -702,12 +740,26 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
if (cudaGetLastError() != cudaSuccess)
return GGML_STATUS_FAILED;
}
- const auto* q = attn->src[0];
- const auto* k = attn->src[1];
+ const auto* q = attn->src[0];
+ const auto* k = attn->src[1];
+ auto* dst = nodes[count - 1];
+ float scale = 1.0f, bias = 0.0f;
+ for (int i = plan.index + 1; i < count; ++i) {
+ if (nodes[i]->op == GGML_OP_SCALE) {
+ scale = param_f32(nodes[i], 0);
+ bias = param_f32(nodes[i], 1);
+ }
+ }
sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
- cudaError_t error = cudaErrorInvalidValue;
+ if (plan.direct_output) {
+ p.output_f32 = static_cast<float*>(dst->data);
+ p.output_scale = scale;
+ p.output_bias = bias;
+ p.scaled_output = dst->op == GGML_OP_SCALE;
+ }
+ cudaError_t error = cudaErrorInvalidValue;
switch (padded_dim(p.dim)) {
case 64:
error = attention_launch<64>(context, p, q->ne[3]);
@@ -733,14 +785,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
}
if (error != cudaSuccess)
return GGML_STATUS_FAILED;
- float scale = 1.0f, bias = 0.0f;
- for (int i = plan.index + 1; i < count; ++i) {
- if (nodes[i]->op == GGML_OP_SCALE) {
- scale = param_f32(nodes[i], 0);
- bias = param_f32(nodes[i], 1);
- }
- }
- auto* dst = nodes[count - 1];
+ if (plan.direct_output)
+ return GGML_STATUS_SUCCESS;
int convrot_group;
if (sd_ggml_convrot_describe(dst, &convrot_group)) {
const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4;
diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu
index 21b6134..9e43d5c 100644
--- a/src/core/cuda/i8_gated_residual.cu
+++ b/src/core/cuda/i8_gated_residual.cu
@@ -229,6 +229,7 @@ namespace sd_i8_gated {
size_t kernel_workspace = 0;
size_t output_offset = 0;
int occupancy = 0;
+ bool output_scratch = true;
ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
};
template <class Types>
@@ -397,24 +398,50 @@ namespace sd_i8_gated {
return configure(describe(nodes, count), plan) ? count : 0;
#endif
}
+ static uintptr_t address(const ggml_tensor* tensor) {
+ size_t offset = 0;
+ while (tensor && !tensor->data && tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
+ }
+ static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ if (!b)
+ return false;
+ const auto x = address(a), y = address(b);
+ if (!x || !y)
+ return false;
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+ static bool needs_scratch(const ggml_tensor* dst, const Operation& op) {
+ if (address(dst) % 16 != 0)
+ return true;
+ for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
+ if (input == op.residual && !op.prefix && address(input) == address(dst))
+ continue;
+ if (overlaps(dst, input))
+ return true;
+ }
+ return false;
+ }
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
auto* plan = new Plan;
- *state = plan;
- plan->device = device_info(ggml_backend_get_device(backend));
- if (!configure(describe(nodes, count), *plan))
+ *state = plan;
+ plan->device = device_info(ggml_backend_get_device(backend));
+ const auto op = describe(nodes, count);
+ if (!configure(op, *plan))
return GGML_STATUS_FAILED;
- *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
+ // Fusion dependencies keep external inputs live until this output is allocated.
+ const auto* dst = nodes[count - 1];
+ const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+ plan->output_scratch = !fresh_output && !(dst->data && !needs_scratch(dst, op));
+ *workspace = plan->output_offset + (plan->output_scratch ? ggml_nbytes(dst) : 0);
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
delete static_cast<Plan*>(state);
}
- static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
- if (!b)
- return false;
- const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
- return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
- }
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
auto& plan = *static_cast<Plan*>(state);
const auto op = describe(nodes, count);
@@ -424,14 +451,11 @@ namespace sd_i8_gated {
ggml_ext_cuda_context context;
if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
return GGML_STATUS_FAILED;
- auto stream = static_cast<cudaStream_t>(context.stream);
- auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
- bool alias = reinterpret_cast<uintptr_t>(dst->data) % 16 != 0;
- for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
- if (input == op.residual && !op.prefix && input->data == dst->data)
- continue;
- alias |= overlaps(dst, input);
- }
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ auto* scratch = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+ const bool alias = needs_scratch(dst, op);
+ if (alias && !plan.output_scratch)
+ return GGML_STATUS_FAILED;
auto* output = alias ? scratch + plan.output_offset : static_cast<char*>(dst->data);
const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
auto status = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);
diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu
index 9d4d92e..c62b95c 100644
--- a/src/core/cuda/i8_gemm.cu
+++ b/src/core/cuda/i8_gemm.cu
@@ -105,6 +105,39 @@ namespace sd_i8_cutlass {
CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
};
+ template <class Base>
+ struct SplitWeightIterator : Base {
+ using Layout = typename Base::Layout;
+ using Pointer = typename Base::Pointer;
+ using TensorCoord = typename Base::TensorCoord;
+ using AccessType = typename Base::AccessType;
+ struct Params : Base::Params {
+ const int8_t* up = nullptr;
+ int64_t k = 0;
+ Params() = default;
+ CUTLASS_HOST_DEVICE Params(const Layout& layout)
+ : Base::Params(layout), k(layout.stride(0)) {}
+ };
+ uintptr_t gate, up;
+ size_t gate_bytes;
+ CUTLASS_DEVICE SplitWeightIterator(const Params& params, Pointer pointer, TensorCoord extent, int thread, TensorCoord offset, const int* indices = nullptr)
+ : Base(params, pointer, extent, thread, offset, indices),
+ gate(reinterpret_cast<uintptr_t>(pointer)),
+ up(reinterpret_cast<uintptr_t>(params.up)),
+ gate_bytes(size_t(params.k) * (extent.column() / 2)) {}
+
+ CUTLASS_DEVICE AccessType* get() const {
+ auto* address = Base::get();
+ if (!address)
+ return nullptr;
+ // Keep the packed SwiGLU iterator's column pairing, but redirect its
+ // virtual up half to a separate allocation. Params.k stays full-sized
+ // even when Stream-K restricts this iterator's K extent.
+ const uintptr_t offset = reinterpret_cast<uintptr_t>(address) - gate;
+ return offset < gate_bytes ? address : reinterpret_cast<AccessType*>(up + offset - gate_bytes);
+ }
+ };
+
template <class Map>
struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
@@ -137,14 +170,15 @@ namespace sd_i8_cutlass {
}
};
- static __global__ void pair_scales(const float* src, float* dst, int n) {
+ static __global__ void pair_scales(const float* src, const float* up, float* dst, int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n)
- dst[i] = src[i / 2 + (i % 2) * (n / 2)];
+ dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)];
}
template <class Config, bool StreamK, bool PackedW4 = false>
struct DenseTypes {
+ static constexpr bool kSplit = false;
using Weight = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
using Tile = typename Config::Tile;
using Warp = typename Config::Warp;
@@ -182,8 +216,9 @@ namespace sd_i8_cutlass {
std::conditional_t<PackedW4, cutlass::arch::OpMultiplyAddMixedInputUpcast, cutlass::arch::OpMultiplyAddSaturate>>::GemmKernel;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
- template <class Config, bool StreamK>
+ template <class Config, bool StreamK, bool Split = false>
struct SwigluTypes {
+ static constexpr bool kSplit = Split;
using Weight = int8_t;
using Tile = typename Config::Tile;
using Warp = typename Config::Warp;
@@ -198,9 +233,23 @@ namespace sd_i8_cutlass {
using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+ using BaseMma = typename GemmBase::Mma;
+ using SplitMma = cutlass::gemm::threadblock::MmaMultistage<
+ typename BaseMma::Shape,
+ typename BaseMma::IteratorA,
+ typename BaseMma::SmemIteratorA,
+ BaseMma::kCacheOpA,
+ SplitWeightIterator<typename BaseMma::IteratorB>,
+ typename BaseMma::SmemIteratorB,
+ BaseMma::kCacheOpB,
+ typename BaseMma::ElementC,
+ typename BaseMma::LayoutC,
+ typename BaseMma::Policy,
+ BaseMma::Detail::kStages>;
+ using Mma = std::conditional_t<Split, SplitMma, BaseMma>;
using Kernel = std::conditional_t<StreamK,
- cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
- cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
+ cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<Mma, VisitorEpilogue, Swizzle>,
+ cutlass::gemm::kernel::GemmWithEpilogueVisitor<Mma, VisitorEpilogue, Swizzle>>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
struct Plan {
@@ -208,9 +257,10 @@ namespace sd_i8_cutlass {
int64_t rows, n, k;
bool swiglu;
bool packed_w4;
+ bool split;
size_t kernel_bytes = 0, scratch_offset = 0;
int occupancy = 0;
- ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t) = nullptr;
+ ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr;
};
template <class Types>
static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
@@ -219,10 +269,10 @@ namespace sd_i8_cutlass {
return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast<const typename Types::Weight*>(w), nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
}
template <class Types>
- static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+ static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales) {
using Kernel = typename Types::Kernel;
if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
- reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out)) %
+ reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales)) %
16)
return GGML_STATUS_FAILED;
auto* paired = static_cast<float*>(scratch);
@@ -239,14 +289,19 @@ namespace sd_i8_cutlass {
return GGML_STATUS_FAILED;
}
typename Kernel::Params params(args, p.device.sms, p.occupancy);
+ if constexpr (Types::kSplit) {
+ if (!up || !up_scales)
+ return GGML_STATUS_FAILED;
+ params.params_B.up = up;
+ }
if (params.get_workspace_size() > p.kernel_bytes)
return GGML_STATUS_FAILED;
if (p.swiglu) {
- pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, paired, int(p.n));
+ pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, up_scales, paired, int(p.n));
if (cudaGetLastError() != cudaSuccess)
return GGML_STATUS_FAILED;
}
- if (params.init_workspace(static_cast<char*>(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess)
+ if (params.init_workspace(scratch ? static_cast<char*>(scratch) + p.scratch_offset : nullptr, stream) != cutlass::Status::kSuccess)
return GGML_STATUS_FAILED;
cutlass::arch::synclog_setup();
cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
@@ -277,6 +332,8 @@ namespace sd_i8_cutlass {
static bool configure_tile(Plan& p) {
if (p.packed_w4)
return configure_kernel<DenseTypes<Config, false, true>>(p);
+ if (p.split)
+ return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true, true>>(p) : configure_kernel<SwigluTypes<Config, false, true>>(p);
if (p.swiglu)
return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
@@ -305,7 +362,7 @@ namespace sd_i8_cutlass {
}
}
struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
-SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) {
+static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4, bool split) {
#ifdef SD_CUDA_FORCE_CUBLAS
return nullptr;
#else
@@ -320,6 +377,7 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
p->k = k;
p->swiglu = swiglu;
p->packed_w4 = packed_w4;
+ p->split = split;
if (!sd_i8_cutlass::configure(*p)) {
delete p;
return nullptr;
@@ -328,8 +386,19 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
return p;
#endif
}
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) {
+ return prepare_gemm(device, rows, n, k, swiglu, workspace, packed_w4, false);
+}
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace) {
+ if (n > INT_MAX / 2 || k > INT_MAX / 16384)
+ return nullptr;
+ return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
+}
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
- return plan->launch(*plan, x, w, xs, ws, out, scratch, stream);
+ return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr);
+}
+ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) {
+ return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales);
}
void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
delete plan;
diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h
index 5e4ac9f..dc0d15d 100644
--- a/src/core/cuda/i8_gemm.h
+++ b/src/core/cuda/i8_gemm.h
@@ -5,5 +5,8 @@
struct SDI8CudaGemmPlan;
SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false);
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
+// n is the width of each branch; gate/up weights and scales remain separate.
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace);
+ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream);
void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan);
#endif
diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu
index 60f36bb..93d8f7f 100644
--- a/src/core/cuda/int8.cu
+++ b/src/core/cuda/int8.cu
@@ -445,6 +445,13 @@ namespace SD_CUDA_NAMESPACE {
static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
return disabled;
}
+#ifdef SD_CUDA_I8_CUTLASS
+ static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) {
+ return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 &&
+ aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
+ (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16)));
+ }
+#endif
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
@@ -466,8 +473,21 @@ namespace SD_CUDA_NAMESPACE {
const auto* quant = ggml_graph_node(g, first + 2);
int group;
if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) &&
- size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
+ size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) {
+#ifdef SD_CUDA_I8_CUTLASS
+ if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) {
+ size_t workspace = 0;
+ auto* gemm = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], true, &workspace);
+ const bool fused_epilogue = gemm != nullptr;
+ sd_cuda_i8_gemm_free(gemm);
+ // ConvRot is a separate launch either way. Keep its input in the
+ // tensor arena so other operations can reuse this large allocation.
+ if (fused_epilogue)
+ return 2;
+ }
+#endif
return 3;
+ }
}
return 2;
}
@@ -511,9 +531,7 @@ namespace SD_CUDA_NAMESPACE {
#ifdef SD_CUDA_I8_CUTLASS
size_t gemm_bytes = 0;
// Retain cuBLAS when I32 accumulation can overflow.
- if (op.weight_scale && !op.bias && p->k <= INT_MAX / 16384 &&
- aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
- (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) &&
+ if (cutlass_compatible(op, dst) &&
(p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
@@ -530,8 +548,12 @@ namespace SD_CUDA_NAMESPACE {
const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
p->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
: 0);
- p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
- *workspace = total;
+ // Custom/concat outputs and packed GLU cannot reuse their input storage.
+ const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src &&
+ (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT ||
+ (dst->op == GGML_OP_GLU && !dst->src[1]));
+ p->output = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
+ *workspace = total;
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
@@ -548,6 +570,99 @@ namespace SD_CUDA_NAMESPACE {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
+#ifdef SD_CUDA_I8_CUTLASS
+ static bool split_describe(ggml_backend_dev_t device, ggml_tensor* const* nodes, SDI8Operation& gate, SDI8Operation& up) {
+ const auto* glu = nodes[2];
+ if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || glu->op_params[1] != 0 ||
+ !glu->src[0] || !glu->src[1] || glu->src[0] == glu->src[1] || glu->view_src ||
+ glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || !aligned(glu, 16) ||
+ !((nodes[0] == glu->src[0] && nodes[1] == glu->src[1]) || (nodes[1] == glu->src[0] && nodes[0] == glu->src[1])))
+ return false;
+ for (const auto* mat : {glu->src[0], glu->src[1]})
+ if (mat->view_src || !supports(device, mat, nullptr) || !ggml_are_same_shape(mat, glu))
+ return false;
+ if (!sd_ggml_i8_describe(glu->src[0], &gate) || !sd_ggml_i8_describe(glu->src[1], &up) ||
+ gate.input != up.input || gate.input_scale != up.input_scale || gate.packed != up.packed ||
+ gate.input->type != GGML_TYPE_I8 || gate.bias || up.bias || !gate.weight_scale || !up.weight_scale ||
+ !ggml_are_same_shape(gate.weight, up.weight))
+ return false;
+ for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+ if (!aligned(t, 16))
+ return false;
+ return true;
+ }
+ static int split_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+ auto* g = const_cast<ggml_cgraph*>(graph);
+ if (fusion_disabled() || first + 2 >= ggml_graph_n_nodes(g))
+ return 0;
+ ggml_tensor* nodes[] = {ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)};
+ SDI8Operation gate, up;
+ if (!split_describe(device, nodes, gate, up))
+ return 0;
+ size_t workspace = 0;
+ auto* gemm = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(nodes[2]), nodes[2]->ne[0], gate.weight->ne[0], &workspace);
+ const bool supported = gemm != nullptr;
+ sd_cuda_i8_gemm_free(gemm);
+ return supported ? 3 : 0;
+ }
+ struct SplitPlan {
+ SDI8CudaGemmPlan* gemm = nullptr;
+ size_t output = SIZE_MAX;
+ int device;
+ };
+ static ggml_status split_prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+ auto* p = new SplitPlan;
+ *state = p;
+ auto device = ggml_backend_get_device(backend);
+ p->device = cuda_device_info(device).device;
+ SDI8Operation gate, up;
+ if (count != 3 || !split_describe(device, nodes, gate, up))
+ return GGML_STATUS_FAILED;
+ auto* dst = nodes[2];
+ size_t bytes = 0;
+ p->gemm = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(dst), dst->ne[0], gate.weight->ne[0], &bytes);
+ if (!p->gemm)
+ return GGML_STATUS_FAILED;
+ size_t total = (bytes + 255) / 256 * 256;
+ if (dst->data || dst->buffer || dst->view_src)
+ p->output = reserve(total, ggml_nbytes(dst));
+ *workspace = total;
+ return GGML_STATUS_SUCCESS;
+ }
+ static void split_destroy(void* state, void*) {
+ auto* p = static_cast<SplitPlan*>(state);
+ if (p)
+ sd_cuda_i8_gemm_free(p->gemm);
+ delete p;
+ }
+ static ggml_status split_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) {
+ auto& p = *static_cast<SplitPlan*>(state);
+ auto* dst = nodes[2];
+ SDI8Operation gate, up;
+ if (!sd_ggml_i8_describe(dst->src[0], &gate) || !sd_ggml_i8_describe(dst->src[1], &up))
+ return GGML_STATUS_FAILED;
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context) || context.device != p.device)
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<CUDAStream>(context.stream);
+ auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ bool alias = false;
+ for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale})
+ alias |= overlaps(dst, t);
+ if (alias && p.output == SIZE_MAX)
+ return GGML_STATUS_FAILED;
+ auto* output = alias ? reinterpret_cast<float*>(scratch + p.output) : static_cast<float*>(dst->data);
+ const auto* input = static_cast<const int8_t*>(gate.input->data);
+ const int64_t padded = (ggml_nrows(dst) + 3) / 4 * 4;
+ const auto* scales = gate.packed ? reinterpret_cast<const float*>(input + gate.weight->ne[0] * padded) : static_cast<const float*>(gate.input_scale->data);
+ auto status = sd_cuda_i8_gemm_execute_split(p.gemm, input, static_cast<const int8_t*>(gate.weight->data), static_cast<const int8_t*>(up.weight->data),
+ scales, static_cast<const float*>(gate.weight_scale->data), static_cast<const float*>(up.weight_scale->data),
+ output, scratch, stream);
+ if (status != GGML_STATUS_SUCCESS || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success))
+ return GGML_STATUS_FAILED;
+ return GGML_STATUS_SUCCESS;
+ }
+#endif
static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
if (p.padded > p.rows) {
auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
@@ -642,9 +757,10 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
return GGML_STATUS_FAILED;
- auto stream = static_cast<CUDAStream>(context.stream);
- auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
- auto* dst = nodes[count - 1];
+ auto stream = static_cast<CUDAStream>(context.stream);
+ auto* scratch = execution->workspace ? static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace)) : nullptr;
+ auto scratch_at = [scratch](size_t offset) { return scratch ? scratch + offset : nullptr; };
+ auto* dst = nodes[count - 1];
SDI8Operation op;
const bool mat = sd_ggml_i8_describe(nodes[0], &op);
bool alias = false;
@@ -659,20 +775,22 @@ namespace SD_CUDA_NAMESPACE {
// CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes.
if (p.gemm && p.kind == SWIGLU_QUANTIZE)
alias = false;
- void* output = alias ? scratch + p.output : dst->data;
+ if (alias && p.output == SIZE_MAX)
+ return GGML_STATUS_FAILED;
+ void* output = alias ? scratch_at(p.output) : dst->data;
CUDAError error = cuda_success;
if (!mat) {
auto* quant = static_cast<int8_t*>(output);
auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
- error = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch + p.partial));
+ error = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch_at(p.partial)));
} else {
- auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
+ auto* quant = reinterpret_cast<int8_t*>(scratch_at(p.quant));
const float* scales;
if (op.input->type == GGML_TYPE_F32) {
- auto* temporary_scales = reinterpret_cast<float*>(scratch + p.scales);
+ auto* temporary_scales = reinterpret_cast<float*>(scratch_at(p.scales));
scales = temporary_scales;
if (op.group_size)
- error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
+ error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch_at(p.partial)));
else {
if (p.padded > p.rows)
error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
@@ -696,12 +814,12 @@ namespace SD_CUDA_NAMESPACE {
const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
auto* gemm_output = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
- auto* accum = reinterpret_cast<int32_t*>(scratch + p.accum);
+ auto* accum = reinterpret_cast<int32_t*>(scratch_at(p.accum));
const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
#ifdef SD_CUDA_I8_CUTLASS
if (p.gemm) {
auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
- intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch + p.gemm_offset, stream);
+ intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream);
if (status != GGML_STATUS_SUCCESS)
return status;
} else
@@ -794,6 +912,17 @@ namespace SD_CUDA_NAMESPACE {
desc.custom_op = sd_ggml_convrot_op();
desc.supports = sd_i8_cuda::convrot_supports;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#ifdef SD_CUDA_I8_CUTLASS
+ desc.name = "sd.linear_i8.swiglu_split.cuda.v1";
+ desc.custom_op = sd_ggml_i8_op();
+ desc.priority = 260;
+ desc.supports = nullptr;
+ desc.match = sd_i8_cuda::split_match;
+ desc.prepare = sd_i8_cuda::split_prepare;
+ desc.execute = sd_i8_cuda::split_execute;
+ desc.destroy = sd_i8_cuda::split_destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+#endif
}
}
diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu
index 9f4192b..02196ff 100644
--- a/src/core/cuda/rms_affine.cu
+++ b/src/core/cuda/rms_affine.cu
@@ -169,15 +169,35 @@ namespace SD_CUDA_NAMESPACE {
}
}
+ static uintptr_t address(const ggml_tensor* tensor) {
+ size_t offset = 0;
+ while (tensor && !tensor->data && tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return tensor && tensor->data ? reinterpret_cast<uintptr_t>(tensor->data) + offset : 0;
+ }
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
if (!b)
return false;
- const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ const auto x = address(a), y = address(b);
+ if (!x || !y)
+ return false;
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
+ static bool needs_scratch(const Operation& op) {
+ const auto* dst = op.output;
+ return (overlaps(dst, op.input) && address(dst) != address(op.input)) ||
+ (overlaps(dst, op.residual) && address(dst) != address(op.residual)) ||
+ overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
+ }
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
- *workspace = ggml_nbytes(nodes[count - 1]);
+ const auto op = describe(nodes, count);
+ const auto* dst = op.output;
+ const bool fresh = !dst->data && !dst->buffer && !dst->view_src;
+ // Fresh outputs may reuse the residual elementwise; external inputs stay live.
+ *workspace = fresh || (dst->data && !needs_scratch(op)) ? 0 : ggml_nbytes(dst);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
@@ -206,11 +226,11 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
- const auto op = describe(nodes, count);
- auto* dst = op.output;
- const bool alias = (overlaps(dst, op.input) && dst->data != op.input->data) ||
- (overlaps(dst, op.residual) && dst->data != op.residual->data) ||
- overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
+ const auto op = describe(nodes, count);
+ auto* dst = op.output;
+ const bool alias = needs_scratch(op);
+ if (alias && !execution->workspace)
+ return GGML_STATUS_FAILED;
auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
const auto stream = static_cast<CUDAStream>(context.stream);
CUDAError error = cuda_invalid_value;
diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu
index 99cc824..f3ab216 100644
--- a/src/core/cuda/rope.cu
+++ b/src/core/cuda/rope.cu
@@ -277,7 +277,10 @@ namespace SD_CUDA_NAMESPACE {
}
static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
- *workspace = ggml_nbytes(nodes[count - 1]);
+ const auto* dst = nodes[count - 1];
+ // Owning casts are fresh outputs; copies into existing storage retain view_src.
+ const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
+ *workspace = fresh_output ? 0 : ggml_nbytes(dst);
return GGML_STATUS_SUCCESS;
}