mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-09 18:36:33 +00:00
perf: reduce fused compute buffers (#2290)
This commit is contained in:
1 parent
9234ad822a
commit
5f8c19f98d
11 files changed
+425
-87
No files matched your search
+34
-5
@@ -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.
|
||||
|
||||
+1
-1
Submodule ggml updated: 135ce1e0e3...3fc478e619.
@@ -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);
|
||||
|
||||
@@ -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*) {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 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))
|
||||
return GGML_STATUS_FAILED;
|
||||
*workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void destroy(void* state, void*) {
|
||||
delete static_cast<Plan*>(state);
|
||||
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 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));
|
||||
const auto op = describe(nodes, count);
|
||||
if (!configure(op, *plan))
|
||||
return GGML_STATUS_FAILED;
|
||||
// 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 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);
|
||||
|
||||
+81
-12
@@ -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;
|
||||
|
||||
@@ -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
|
||||
+145
-16
@@ -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
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user