DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-12 11:56:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

revert: restore public backend extension integration

Browse Source
This commit is contained in:
leejet committed 2026-11-01 04:09:12 +08:00
1 parent a013932674
commit 75e8157b01
26 files changed
+590 -358
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cpu/int8.cpp
+12 -3
View File
@@ -101,18 +101,24 @@ namespace sd_cpu {
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
}
static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
auto* dst = nodes[0];
I8Execution e = {};
if (!sd_ggml_i8_describe(dst, &e.op))
return GGML_STATUS_FAILED;
e.rows = ggml_nrows(dst);
auto run = parallel(context->backend);
if (!run)
return GGML_STATUS_FAILED;
char* scratch = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
const size_t quant_size = quant_workspace(dst, e.op);
if (e.op.input->type == GGML_TYPE_F32) {
e.quant = reinterpret_cast<int8_t*>(scratch);
e.scales = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
auto status = parallel(context->backend, dst, quantize, &e);
auto status = run(context->backend, dst, quantize, &e);
if (status != GGML_STATUS_SUCCESS)
return status;
} else {
@@ -123,7 +129,7 @@ namespace sd_cpu {
for (auto* input : dst->src)
alias |= overlaps(dst, input);
e.output = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
auto status = parallel(context->backend, dst, multiply, &e);
auto status = run(context->backend, dst, multiply, &e);
if (status == GGML_STATUS_SUCCESS && alias)
std::memcpy(dst->data, e.output, ggml_nbytes(dst));
return status;
@@ -133,13 +139,16 @@ namespace sd_cpu {
I8Execution e = {};
if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
return GGML_STATUS_FAILED;
auto run = parallel(context->backend);
if (!run)
return GGML_STATUS_FAILED;
e.op.input = dst->src[0];
e.rows = ggml_nrows(e.op.input);
bool alias = overlaps(dst, e.op.input);
e.quant = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
e.scales = reinterpret_cast<float*>(e.quant + k * padded);
auto status = parallel(context->backend, dst, quantize, &e);
auto status = run(context->backend, dst, quantize, &e);
if (status == GGML_STATUS_SUCCESS) {
std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
if (alias)
src/core/cpu/registration.cpp
-21
View File
@@ -1,29 +1,8 @@
#include "core/cpu/registration.h"
#ifndef SD_USE_UPSTREAM_GGML
#include <cstring>
#include "ggml-backend-impl.h"
#include "ggml-cpu.h"
#include "ggml-impl.h"
namespace sd_cpu {
ggml_status parallel(ggml_backend_t backend, ggml_tensor* output, ggml_custom_op_t callback, void* user_data) {
if (!ggml_backend_is_cpu(backend) || !output || !callback) {
return GGML_STATUS_FAILED;
}
ggml_tensor node = *output;
ggml_custom_op_params params = {callback, GGML_N_TASKS_MAX, user_data};
node.op = GGML_OP_CUSTOM;
node.flags |= GGML_TENSOR_FLAG_COMPUTE;
memcpy(node.op_params, &params, sizeof(params));
ggml_tensor* nodes[] = {&node};
ggml_cgraph graph = {};
graph.n_nodes = 1;
graph.nodes = nodes;
// Reuse the configured thread pool without dispatching this synthetic node through extensions.
return backend->iface.graph_compute(backend, &graph);
}
void register_extensions(ggml_backend_dev_t device) {
register_i8(device);
register_rope(device);
src/core/cpu/registration.h
-1
View File
@@ -5,7 +5,6 @@
#ifndef SD_USE_UPSTREAM_GGML
namespace sd_cpu {
ggml_status parallel(ggml_backend_t backend, ggml_tensor* output, ggml_custom_op_t callback, void* user_data);
void register_extensions(ggml_backend_dev_t device);
void register_rope(ggml_backend_dev_t device);
void register_i8(ggml_backend_dev_t device);
src/core/cpu/rope.cpp
+5 -1
View File
@@ -103,12 +103,16 @@ namespace sd_cpu {
}
static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
auto device = ggml_backend_get_device(context->backend);
auto reg = ggml_backend_dev_backend_reg(device);
auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
GGML_ASSERT(run);
RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
if (count == 3) {
execution.norm = nodes[0];
execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
}
return parallel(context->backend, nodes[count - 1], rope_compute, &execution);
return run(context->backend, nodes[count - 1], rope_compute, &execution);
}
void register_rope(ggml_backend_dev_t device) {
src/core/cuda/CMakeLists.txt
+1 -1
View File
@@ -11,7 +11,7 @@ function(sd_cuda_sources VENDOR LANGUAGE)
endforeach()
set(TARGET "sd-${VENDOR}")
add_library(${TARGET} OBJECT ${SOURCES})
target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src" "${SD_GGML_CUDA_INCLUDE_DIR}" "${sd_ggml_private_include}")
target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
target_link_libraries(${TARGET} PRIVATE ggml)
set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_compile_features(${TARGET} PRIVATE cxx_std_17)
src/core/cuda/affine.cu
+6 -6
View File
@@ -1,8 +1,8 @@
#include <climits>
#include <cstdlib>
#include <cstring>
#include "backend.h"
#include "ggml-ext.h"
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_affine_cuda {
@@ -129,7 +129,7 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
auto stream = static_cast<CUDAStream>(context.stream);
const auto* node = nodes[0];
const auto* a = node->src[0];
const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
@@ -154,10 +154,10 @@ namespace SD_CUDA_NAMESPACE {
scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
}
auto error = cudaGetLastError();
if (error == cudaSuccess && alias)
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
auto error = cuda_last_error();
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
}
src/core/cuda/backend.h
-41
View File
@@ -1,41 +0,0 @@
#ifndef __SD_CORE_CUDA_BACKEND_H__
#define __SD_CORE_CUDA_BACKEND_H__
#include "ggml-ext-gpu.h"
#include "ggml-impl.h"
#if defined(SD_CUDA_VENDOR_HIP)
#include "vendors/hip.h"
#define SD_CUDA_NAMESPACE sd_hip
#define SD_CUDA_NAME "hip"
#elif defined(SD_CUDA_VENDOR_MUSA)
#include "vendors/musa.h"
#define SD_CUDA_NAMESPACE sd_musa
#define SD_CUDA_NAME "musa"
#else
#include "vendors/cuda.h"
#define SD_CUDA_NAMESPACE sd_cuda
#define SD_CUDA_NAME "cuda"
#endif
namespace SD_CUDA_NAMESPACE {
#if defined(SD_CUDA_VENDOR_HIP)
static constexpr auto cuda_kind = GGML_EXT_GPU_HIP;
#elif defined(SD_CUDA_VENDOR_MUSA)
static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA;
#else
static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA;
#endif
static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
return query && query(backend, context) && context->kind == cuda_kind;
}
static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
ggml_ext_gpu_device_info info = {};
auto query = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
if (!query || !query(device, &info) || info.kind != cuda_kind)
return {cuda_kind, -1, 0, 0, false};
return info;
}
}
#endif
src/core/cuda/convrot.cu
+19 -15
View File
@@ -3,9 +3,9 @@
#include <cstdlib>
#include <cstring>
#include <string>
#include "backend.h"
#include "convrot_kernel.h"
#include "core/ggml_extensions.h"
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_convrot {
@@ -14,6 +14,11 @@ 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;
}
static float param_f32(const ggml_tensor* tensor, int index) {
float value;
std::memcpy(&value, tensor->op_params + index, sizeof(value));
return value;
}
static bool f32(const ggml_tensor* tensor) {
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
}
@@ -52,12 +57,12 @@ namespace SD_CUDA_NAMESPACE {
op.up_stride = op.k;
if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
op.kind = SCALE;
op.scale = ggml_get_op_params_f32(first, 0);
op.bias = ggml_get_op_params_f32(first, 1);
op.scale = param_f32(first, 0);
op.bias = param_f32(first, 1);
} else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
op.kind = NORM;
op.modulation = nodes[1]->src[1];
op.eps = ggml_get_op_params_f32(first, 0);
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.modulation, op.k))
return false;
} else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
@@ -68,7 +73,7 @@ namespace SD_CUDA_NAMESPACE {
op.weight = nodes[1]->src[1];
op.modulation = nodes[2]->src[1];
op.shift = nodes[4]->src[1];
op.eps = ggml_get_op_params_f32(first, 0);
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
return false;
} else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
@@ -206,16 +211,15 @@ namespace SD_CUDA_NAMESPACE {
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
static cudaError_t launch(Plan& p, const Operation& op, int8_t* output, cudaStream_t stream) {
static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
constexpr bool warp = Scaled || RMS;
constexpr int extra = Normalize ? 384 : RMS ? 256
: 128;
if (!p.configured) {
#if !defined(SD_CUDA_VENDOR_HIP)
auto status = cudaFuncSetAttribute(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - extra));
if (status != cudaSuccess)
auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
int(p.device.shared_memory - extra));
if (status != cuda_success)
return status;
#endif
p.configured = true;
}
const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
@@ -228,7 +232,7 @@ namespace SD_CUDA_NAMESPACE {
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
return cudaGetLastError();
return cuda_last_error();
}
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);
@@ -237,15 +241,15 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
auto stream = static_cast<CUDAStream>(context.stream);
bool alias = false;
for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
alias |= overlaps(dst, input);
auto* output = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
const int64_t padded = (op.rows + 3) / 4 * 4;
if (padded > op.rows && cudaMemsetAsync(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cudaSuccess)
if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
return GGML_STATUS_FAILED;
cudaError_t error;
CUDAError error;
switch (op.kind) {
case SCALE:
error = launch<false, false, true>(plan, op, output, stream);
@@ -266,7 +270,7 @@ namespace SD_CUDA_NAMESPACE {
default:
return GGML_STATUS_FAILED;
}
if (error != cudaSuccess || (alias && cudaMemcpyAsync(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), cudaMemcpyDeviceToDevice, stream) != cudaSuccess))
if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
return GGML_STATUS_FAILED;
return GGML_STATUS_SUCCESS;
}
src/core/cuda/convrot_kernel.h
+14 -14
View File
@@ -1,28 +1,28 @@
#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
#define __SD_CORE_CUDA_CONVROT_KERNEL_H__
#include <cstdint>
#include "backend.h"
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_convrot {
static __device__ __forceinline__ float warp_reduce_sum(float x) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
x += __shfl_xor_sync(0xffffffff, x, offset, 32);
x += cuda_shuffle_xor(x, offset, 32);
return x;
}
static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
x.x += __shfl_xor_sync(0xffffffff, x.x, offset, 32);
x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32);
x.x += cuda_shuffle_xor(x.x, offset, 32);
x.y += cuda_shuffle_xor(x.y, offset, 32);
}
return x;
}
template <int Width = 32>
template <int Width>
static __device__ __forceinline__ float warp_reduce_max(float x) {
#pragma unroll
for (int offset = Width / 2; offset > 0; offset >>= 1)
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset, Width));
x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
return x;
}
static __device__ float2 block_sum(float2 x, float2* shared) {
@@ -130,19 +130,19 @@ namespace SD_CUDA_NAMESPACE {
const int base = lane & ~(3 * stride);
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float a = __shfl_sync(0xffffffff, v[j], base, 32);
const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
const float a = cuda_shuffle(v[j], base, 32);
const float b = cuda_shuffle(v[j], base + stride, 32);
const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
}
}
#pragma unroll
for (int j = 0; j < 8; j += 2) {
const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
const float a = cuda_shuffle(v[j], lane % 16, 32);
const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
}
src/core/cuda/cute_attention.cu
+44 -15
View File
@@ -1,4 +1,3 @@
#include "backend.h"
/*
* Adapted from FlashAttention v2.8.3:
* https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
@@ -461,23 +460,51 @@ static __global__ void copy_from_half(const half* src, float* dst, int64_t count
dst[i] = scaled ? __fmaf_rn(__half2float(src[i]), scale, bias) : __half2float(src[i]);
}
static float param_f32(const ggml_tensor* node, int index) {
float value;
std::memcpy(&value, reinterpret_cast<const char*>(node->op_params) + index * sizeof(value), sizeof(value));
return value;
}
static int padded_dim(int64_t dim) {
return dim == 80 || dim == 96 || dim == 112 ? int(dim) : int((dim + 63) / 64 * 64);
}
struct CUDADevice {
int ordinal;
cudaDeviceProp properties;
};
static const CUDADevice* cuda_device(ggml_backend_dev_t device) {
static std::mutex mutex;
static std::unordered_map<ggml_backend_dev_t, CUDADevice> devices;
std::lock_guard<std::mutex> lock(mutex);
auto found = devices.find(device);
if (found != devices.end())
return &found->second;
auto reg = ggml_backend_dev_backend_reg(device);
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
if (!query)
return nullptr;
CUDADevice info = {};
info.ordinal = query(device);
if (info.ordinal < 0 || cudaGetDeviceProperties(&info.properties, info.ordinal) != cudaSuccess)
return nullptr;
return &devices.emplace(device, info).first->second;
}
static bool attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
static const bool disabled = std::getenv("GGML_CUDA_DISABLE_CUTE_FATTN") || std::getenv("GGML_CUDA_DISABLE_CUTLASS_FATTN");
const auto info = sd_cuda::cuda_device_info(device);
const auto* info = cuda_device(device);
const auto* q = node->src[0];
const auto* k = node->src[1];
const auto* v = node->src[2];
if (disabled || info.device < 0 || info.cc < 800 || !q || !k || !v ||
if (disabled || !info || info->properties.major < 8 || !q || !k || !v ||
q->type != GGML_TYPE_F32 || node->type != GGML_TYPE_F32 ||
q->ne[0] < 32 || q->ne[0] > 256 || q->ne[0] % 8 ||
size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info.shared_memory ||
size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info->properties.sharedMemPerBlockOptin ||
q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
node->src[3] || node->src[4] || ggml_get_op_params_f32(node, 1) != 0.0f || ggml_get_op_params_f32(node, 2) != 0.0f ||
node->src[3] || node->src[4] || param_f32(node, 1) != 0.0f || param_f32(node, 2) != 0.0f ||
q->ne[2] % k->ne[2] || q->ne[3] != k->ne[3] || !ggml_are_same_shape(k, v) ||
q->nb[0] != sizeof(float) || !ggml_is_contiguous(node) ||
(k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) ||
@@ -535,7 +562,7 @@ static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, in
scaled = true;
} else if (sd_ggml_convrot_describe(node, &convrot_group)) {
if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
size_t(previous->ne[0]) * sizeof(float) + 128 <= sd_cuda::cuda_device_info(device).shared_memory)
size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
return count + 1;
break;
} else
@@ -617,7 +644,7 @@ static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes,
}
template <int D, bool FullDim>
static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
auto kernel = sd_fattn_cute::kernel<64, 64, D, FullDim>;
constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
static std::mutex mutex;
@@ -636,7 +663,7 @@ static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const s
}
template <int D>
static cudaError_t attention_launch(const ggml_ext_gpu_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
}
@@ -644,8 +671,10 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
const auto& plan = *static_cast<AttentionPlan*>(state);
const auto* attn = nodes[plan.index];
auto device = ggml_backend_get_device(execution->backend);
ggml_ext_gpu_context context;
if (!sd_cuda::cuda_context(execution->backend, &context))
auto reg = ggml_backend_dev_backend_reg(device);
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
ggml_ext_cuda_context context;
if (!query || !query(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
auto* base = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
@@ -655,7 +684,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
const half* inputs[3] = {work[0], work[1], work[2]};
for (int i = 0; i < 3; ++i) {
const auto* src = attn->src[i];
const float scale = i == 0 ? ggml_get_op_params_f32(attn, 0) : 1.0f;
const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
if (i == 0 && plan.index == 5) {
if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
return GGML_STATUS_FAILED;
@@ -706,8 +735,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
float scale = 1.0f, bias = 0.0f;
for (int i = plan.index + 1; i < count; ++i) {
if (nodes[i]->op == GGML_OP_SCALE) {
scale = ggml_get_op_params_f32(nodes[i], 0);
bias = ggml_get_op_params_f32(nodes[i], 1);
scale = param_f32(nodes[i], 0);
bias = param_f32(nodes[i], 1);
}
}
auto* dst = nodes[count - 1];
@@ -718,7 +747,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
return GGML_STATUS_FAILED;
if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
int(sd_cuda::cuda_device_info(device).shared_memory - 128)) != cudaSuccess)
int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
return GGML_STATUS_FAILED;
quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
@@ -733,7 +762,7 @@ static void attention_destroy(void* state, void*) {
}
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
if (sd_cuda::cuda_device_info(device).device < 0)
if (!cuda_device(device))
return;
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
src/core/cuda/i8_gated_residual.cu
+118 -27
View File
@@ -7,11 +7,71 @@
#include <cstdlib>
#include <cstring>
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
#include "backend.h"
#include <mutex>
#include <unordered_map>
#include "core/ggml_extensions.h"
#include "i8_gemm.cuh"
namespace sd_i8_gated {
struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
using Base::Base;
CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
const auto shape = tiled_shape();
if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
// Column bands retain weight tiles in Ada's L2 across successive row tiles.
constexpr int band_width = 16;
if (index >= shape.m() * shape.n()) {
return {shape.m(), shape.n(), get_batch_idx()};
}
const int band = index / (shape.m() * band_width);
const int width = min(band_width, shape.n() - band * band_width);
const int within = index - band * shape.m() * band_width;
return {within / width, band * band_width + within % width, get_batch_idx()};
}
#endif
return Base::get_tile_offset(index);
}
};
template <int M, int N, int WM>
struct TileConfig {
using Tile = cutlass::gemm::GemmShape<M, N, 64>;
using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
};
enum Tile {
TILE_128_256,
TILE_128_128,
TILE_64_128,
TILE_16_128,
};
static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
}
static Tile select_tile(int64_t rows, int64_t n, int nsm) {
if (rows < 32) {
return TILE_16_128;
}
const int64_t columns = (n + 127) / 128;
if (((rows + 15) / 16) * columns <= nsm) {
return TILE_16_128;
}
if (((rows + 63) / 64) * columns <= nsm) {
return TILE_64_128;
}
const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
const int64_t square_tiles = ((rows + 127) / 128) * columns;
// A narrower tile costs more per output element, but can waste less of the last wave.
const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
}
template <class T>
struct MulRound {
CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
@@ -40,7 +100,7 @@ namespace sd_i8_gated {
struct KernelTypes {
using Tile = typename Config::Tile;
using Warp = typename Config::Warp;
using Swizzle = std::conditional_t<stream_k, sd_i8_cutlass::BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
using Swizzle = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
@@ -84,6 +144,35 @@ namespace sd_i8_gated {
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
struct DeviceInfo {
int ordinal = -1;
int major = 0;
int sms = 0;
int threads_per_sm = 0;
size_t shared = 0;
};
static DeviceInfo device_info(ggml_backend_dev_t device) {
static std::mutex mutex;
static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
std::lock_guard<std::mutex> lock(mutex);
auto found = cache.find(device);
if (found != cache.end())
return found->second;
DeviceInfo info;
auto reg = ggml_backend_dev_backend_reg(device);
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
int ordinal = query ? query(device) : -1;
cudaDeviceProp props = {};
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
cache.emplace(device, info);
return info;
}
static float param_f32(const ggml_tensor* tensor, int index) {
float value;
std::memcpy(&value, tensor->op_params + index, sizeof(value));
return value;
}
static bool aligned(const ggml_tensor* tensor) {
if (!tensor)
return false;
@@ -111,7 +200,7 @@ namespace sd_i8_gated {
op.gate = mul->src[1];
op.residual = add->src[0] == mul ? add->src[1] : add->src[0];
op.prefix = concat ? nodes[count - 1]->src[0] : nullptr;
op.scale = nodes[1]->op == GGML_OP_SCALE ? ggml_get_op_params_f32(nodes[1], 0) : 1.f;
op.scale = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
op.rows = ggml_nrows(nodes[0]);
op.n = nodes[0]->ne[0];
op.k = op.mat.weight->ne[0];
@@ -136,7 +225,7 @@ namespace sd_i8_gated {
return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
}
struct Plan {
ggml_ext_gpu_device_info device;
DeviceInfo device;
size_t kernel_workspace = 0;
size_t output_offset = 0;
int occupancy = 0;
@@ -159,7 +248,7 @@ namespace sd_i8_gated {
return GGML_STATUS_FAILED;
plan.occupancy = occupancy;
}
typename Kernel::Params params(args, plan.device.nsm, plan.occupancy);
typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
return GGML_STATUS_FAILED;
cutlass::arch::synclog_setup();
@@ -170,15 +259,15 @@ namespace sd_i8_gated {
static bool configure(const Operation& op, Plan& plan) {
using Types = KernelTypes<Config, Scaled, StreamK>;
using Kernel = typename Types::Kernel;
if (sizeof(typename Kernel::SharedStorage) > plan.device.shared_memory)
if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
return false;
auto args = arguments<Types>(op, nullptr, true);
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
return false;
size_t workspace = 0;
// Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
for (int occupancy = 1; occupancy <= plan.device.max_threads_per_sm / Kernel::kThreadCount; ++occupancy) {
typename Kernel::Params params(args, plan.device.nsm, occupancy);
for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
typename Kernel::Params params(args, plan.device.sms, occupancy);
const auto grid = params.get_grid_dims();
if (grid.y > 65535 || grid.z > 65535)
return false;
@@ -196,23 +285,23 @@ namespace sd_i8_gated {
return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
}
static bool configure(const Operation& op, Plan& plan) {
auto try_tile = [&](sd_i8_cutlass::Tile tile) {
auto try_tile = [&](Tile tile) {
switch (tile) {
case sd_i8_cutlass::TILE_128_256:
return configure<sd_i8_cutlass::TileConfig<128, 256, 64>>(op, plan);
case sd_i8_cutlass::TILE_128_128:
return configure<sd_i8_cutlass::TileConfig<128, 128, 64>>(op, plan);
case sd_i8_cutlass::TILE_64_128:
return configure<sd_i8_cutlass::TileConfig<64, 128, 32>>(op, plan);
case sd_i8_cutlass::TILE_16_128:
return configure<sd_i8_cutlass::TileConfig<16, 128, 16>>(op, plan);
case TILE_128_256:
return configure<TileConfig<128, 256, 64>>(op, plan);
case TILE_128_128:
return configure<TileConfig<128, 128, 64>>(op, plan);
case TILE_64_128:
return configure<TileConfig<64, 128, 32>>(op, plan);
case TILE_16_128:
return configure<TileConfig<16, 128, 16>>(op, plan);
}
return false;
};
const auto tile = sd_i8_cutlass::select_tile(op.rows, op.n, plan.device.nsm);
const auto tile = select_tile(op.rows, op.n, plan.device.sms);
if (try_tile(tile))
return true;
for (auto fallback : {sd_i8_cutlass::TILE_128_256, sd_i8_cutlass::TILE_128_128, sd_i8_cutlass::TILE_64_128, sd_i8_cutlass::TILE_16_128})
for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
if (fallback != tile && try_tile(fallback))
return true;
return false;
@@ -268,7 +357,7 @@ namespace sd_i8_gated {
if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
return 0;
const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
if (rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
return 0;
for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
@@ -276,7 +365,7 @@ namespace sd_i8_gated {
return 0;
const bool scaled = nodes[1]->op == GGML_OP_SCALE;
int count = scaled ? 4 : 3;
if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || ggml_get_op_params_f32(nodes[1], 1) != 0.f)))
if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
return 0;
const auto* mul = nodes[count - 2];
const auto* add = nodes[count - 1];
@@ -302,8 +391,8 @@ namespace sd_i8_gated {
if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
return 0;
Plan plan;
plan.device = sd_cuda::cuda_device_info(device);
if (plan.device.cc < 800 || plan.device.nsm <= 0)
plan.device = device_info(device);
if (plan.device.major < 8 || plan.device.sms <= 0)
return 0;
return configure(describe(nodes, count), plan) ? count : 0;
#endif
@@ -311,7 +400,7 @@ namespace sd_i8_gated {
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 = sd_cuda::cuda_device_info(ggml_backend_get_device(backend));
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]);
@@ -330,8 +419,10 @@ namespace sd_i8_gated {
auto& plan = *static_cast<Plan*>(state);
const auto op = describe(nodes, count);
auto* dst = nodes[count - 1];
ggml_ext_gpu_context context;
if (!sd_cuda::cuda_context(execution->backend, &context) || context.device != plan.device.device)
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
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));
src/core/cuda/i8_gemm.cu
+94 -10
View File
@@ -5,11 +5,95 @@
#include <climits>
#include <cstdlib>
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
#include "backend.h"
#include "i8_gemm.cuh"
#include <mutex>
#include <unordered_map>
#include "i8_gemm.h"
namespace sd_i8_cutlass {
struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
using Base::Base;
CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
const auto shape = tiled_shape();
if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
// Column bands retain weight tiles in Ada's L2 across successive row tiles.
constexpr int band_width = 16;
if (index >= shape.m() * shape.n()) {
return {shape.m(), shape.n(), get_batch_idx()};
}
const int band = index / (shape.m() * band_width);
const int width = min(band_width, shape.n() - band * band_width);
const int within = index - band * shape.m() * band_width;
return {within / width, band * band_width + within % width, get_batch_idx()};
}
#endif
return Base::get_tile_offset(index);
}
};
template <int M, int N, int WM>
struct TileConfig {
using Tile = cutlass::gemm::GemmShape<M, N, 64>;
using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
};
enum Tile {
TILE_128_256,
TILE_128_128,
TILE_64_128,
TILE_16_128,
};
static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
}
static Tile select_tile(int64_t rows, int64_t n, int nsm) {
if (rows < 32) {
return TILE_16_128;
}
const int64_t columns = (n + 127) / 128;
if (((rows + 15) / 16) * columns <= nsm) {
return TILE_16_128;
}
if (((rows + 63) / 64) * columns <= nsm) {
return TILE_64_128;
}
const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
const int64_t square_tiles = ((rows + 127) / 128) * columns;
// A narrower tile costs more per output element, but can waste less of the last wave.
const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
}
struct DeviceInfo {
int ordinal = -1;
int major = 0;
int sms = 0;
int threads_per_sm = 0;
size_t shared = 0;
};
static DeviceInfo device_info(ggml_backend_dev_t device) {
static std::mutex mutex;
static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
std::lock_guard<std::mutex> lock(mutex);
auto found = cache.find(device);
if (found != cache.end())
return found->second;
DeviceInfo info;
auto reg = ggml_backend_dev_backend_reg(device);
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
int ordinal = query ? query(device) : -1;
cudaDeviceProp props = {};
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
cache.emplace(device, info);
return info;
}
// The iterator pairs gate/up columns while reading the original checkpoint layout.
struct WeightLayout : cutlass::layout::PermuteBase {
int stride, half_n;
@@ -118,7 +202,7 @@ namespace sd_i8_cutlass {
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
struct Plan {
ggml_ext_gpu_device_info device;
DeviceInfo device;
int64_t rows, n, k;
bool swiglu;
size_t kernel_bytes = 0, scratch_offset = 0;
@@ -151,7 +235,7 @@ namespace sd_i8_cutlass {
p.occupancy <= 0)
return GGML_STATUS_FAILED;
}
typename Kernel::Params params(args, p.device.nsm, p.occupancy);
typename Kernel::Params params(args, p.device.sms, p.occupancy);
if (params.get_workspace_size() > p.kernel_bytes)
return GGML_STATUS_FAILED;
if (p.swiglu) {
@@ -168,14 +252,14 @@ namespace sd_i8_cutlass {
template <class Types>
static bool configure_kernel(Plan& p) {
using Kernel = typename Types::Kernel;
if (sizeof(typename Kernel::SharedStorage) > p.device.shared_memory)
if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
return false;
auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
return false;
size_t bytes = 0;
for (int occupancy = 1; occupancy <= p.device.max_threads_per_sm / Kernel::kThreadCount; ++occupancy) {
typename Kernel::Params params(args, p.device.nsm, occupancy);
for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
typename Kernel::Params params(args, p.device.sms, occupancy);
const auto grid = params.get_grid_dims();
if (grid.y > 65535 || grid.z > 65535)
return false;
@@ -206,7 +290,7 @@ namespace sd_i8_cutlass {
}
return false;
};
const auto tile = select_tile(p.rows, p.n, p.device.nsm);
const auto tile = select_tile(p.rows, p.n, p.device.sms);
if (try_tile(tile))
return true;
for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
@@ -220,8 +304,8 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
#ifdef SD_CUDA_FORCE_CUBLAS
return nullptr;
#else
const auto info = sd_cuda::cuda_device_info(device);
if (info.device < 0 || info.cc < 800 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
const auto info = sd_i8_cutlass::device_info(device);
if (info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
return nullptr;
auto* p = new SDI8CudaGemmPlan;
src/core/cuda/i8_gemm.cuh
-70
View File
@@ -1,70 +0,0 @@
#ifndef __SD_CORE_CUDA_I8_GEMM_CUH__
#define __SD_CORE_CUDA_I8_GEMM_CUH__
#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
#include <cstdlib>
namespace sd_i8_cutlass {
struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
using Base::Base;
CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
const auto shape = tiled_shape();
if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
// Column bands retain weight tiles in Ada's L2 across successive row tiles.
constexpr int band_width = 16;
if (index >= shape.m() * shape.n()) {
return {shape.m(), shape.n(), get_batch_idx()};
}
const int band = index / (shape.m() * band_width);
const int width = min(band_width, shape.n() - band * band_width);
const int within = index - band * shape.m() * band_width;
return {within / width, band * band_width + within % width, get_batch_idx()};
}
#endif
return Base::get_tile_offset(index);
}
};
template <int M, int N, int WM>
struct TileConfig {
using Tile = cutlass::gemm::GemmShape<M, N, 64>;
using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
};
enum Tile {
TILE_128_256,
TILE_128_128,
TILE_64_128,
TILE_16_128,
};
static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
}
static Tile select_tile(int64_t rows, int64_t n, int nsm) {
if (rows < 32) {
return TILE_16_128;
}
const int64_t columns = (n + 127) / 128;
if (((rows + 15) / 16) * columns <= nsm) {
return TILE_16_128;
}
if (((rows + 63) / 64) * columns <= nsm) {
return TILE_64_128;
}
const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
const int64_t square_tiles = ((rows + 127) / 128) * columns;
// A narrower tile costs more per output element, but can waste less of the last wave.
const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
}
}
#endif
src/core/cuda/i8_projection.cu
+27 -8
View File
@@ -1,4 +1,3 @@
#include "backend.h"
#include "core/cuda/i8_projection.h"
#include "core/cuda/rope.h"
#include "core/ggml_extensions.h"
@@ -11,7 +10,9 @@
#include <cstdlib>
#include <cstring>
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
#include <mutex>
#include <type_traits>
#include <unordered_map>
static __device__ __forceinline__ float warp_sum8(float value) {
#pragma unroll
@@ -149,16 +150,32 @@ struct ProjectionKernel {
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
static float param_f32(const ggml_tensor* tensor, int index) {
float result;
std::memcpy(&result, tensor->op_params + index, sizeof(result));
return result;
}
static bool device_supports(ggml_backend_dev_t device) {
#ifdef SD_CUDA_FORCE_CUBLAS
(void)device;
return false;
#else
const auto info = sd_cuda::cuda_device_info(device);
static std::mutex mutex;
static std::unordered_map<ggml_backend_dev_t, bool> cache;
std::lock_guard<std::mutex> lock(mutex);
auto found = cache.find(device);
if (found != cache.end())
return found->second;
auto reg = ggml_backend_dev_backend_reg(device);
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
int ordinal = query ? query(device) : -1;
cudaDeviceProp props = {};
size_t shared = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
return info.device >= 0 && info.cc >= 800 && info.shared_memory >= shared;
bool supported = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
cache.emplace(device, supported);
return supported;
#endif
}
@@ -222,7 +239,7 @@ static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat
reinterpret_cast<uintptr_t>(output) % 8)
return cudaErrorInvalidValue;
typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
{{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? ggml_get_op_params_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
{{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
Gemm gemm;
@@ -284,7 +301,7 @@ static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
const auto* rope = node_at(graph, first + 4);
const auto* scale = node_at(graph, first + 5);
const auto* copy = node_at(graph, first + 6);
if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && ggml_get_op_params_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 7;
}
@@ -330,8 +347,10 @@ static ggml_status projection_execute(const ggml_ext_execution_context* executio
const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
const bool joined = concat->op == GGML_OP_CONCAT;
auto* dst = nodes[count - 1];
ggml_ext_gpu_context context;
if (!sd_cuda::cuda_context(execution->backend, &context))
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
ggml_ext_cuda_context context;
if (!query || !query(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
bool alias = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
@@ -340,7 +359,7 @@ static ggml_status projection_execute(const ggml_ext_execution_context* executio
alias |= overlaps(dst, input);
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
float scale = key && count == 7 && !joined ? ggml_get_op_params_f32(nodes[5], 0) : 1.f;
float scale = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
cudaError_t error;
if (dst->type == GGML_TYPE_F32)
error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);
src/core/cuda/int8.cu
+68 -43
View File
@@ -7,17 +7,26 @@
struct SDI8CudaGemmPlan;
#endif
#include "runtime.h"
#if defined(SD_CUDA_VENDOR_HIP)
#include <hipblas/hipblas.h>
#elif !defined(SD_CUDA_VENDOR_MUSA)
#include <cublas_v2.h>
#endif
#include <algorithm>
#include <climits>
#include <mutex>
#include <unordered_map>
#include "backend.h"
namespace SD_CUDA_NAMESPACE {
namespace sd_i8_cuda {
using sd_convrot::convrot_h4_cuda;
using sd_convrot::warp_reduce_max;
template <int Width = 32>
static __device__ __forceinline__ float warp_reduce_max(float value) {
#pragma unroll
for (int offset = Width / 2; offset > 0; offset >>= 1)
value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
return value;
}
static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
const int64_t row = blockIdx.x;
const int tid = threadIdx.x;
@@ -60,6 +69,13 @@ namespace SD_CUDA_NAMESPACE {
dst_row[i] = (int8_t)value;
}
}
static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
a = index == 3 ? -a : a;
b = index == 2 ? -b : b;
c = index == 1 ? -c : c;
d = index == 0 ? -d : d;
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
}
static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
constexpr int group_size = 256;
const int64_t row = blockIdx.x;
@@ -86,19 +102,19 @@ namespace SD_CUDA_NAMESPACE {
const int base = lane & ~(3 * stride);
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float a = __shfl_sync(0xffffffff, v[j], base, 32);
const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
const float a = cuda_shuffle(v[j], base, 32);
const float b = cuda_shuffle(v[j], base + stride, 32);
const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
}
}
#pragma unroll
for (int j = 0; j < 8; j += 2) {
const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
const float a = cuda_shuffle(v[j], lane % 16, 32);
const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
}
@@ -483,10 +499,10 @@ 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);
}
static cudaError_t convrot_launch(const Plan& p, cudaStream_t stream, const float* src, int8_t* dst, float* scales, float* partial) {
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 = cudaMemsetAsync(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
if (error != cudaSuccess)
auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
if (error != cuda_success)
return error;
}
const size_t shared = size_t(p.k) * sizeof(float);
@@ -496,11 +512,9 @@ namespace SD_CUDA_NAMESPACE {
{
std::lock_guard<std::mutex> lock(mutex);
if (!configured[p.device.device]) {
#if !defined(SD_CUDA_VENDOR_HIP)
auto error = cudaFuncSetAttribute(quantize_convrot, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - 128));
if (error != cudaSuccess)
auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
if (error != cuda_success)
return error;
#endif
configured[p.device.device] = true;
}
}
@@ -513,7 +527,7 @@ namespace SD_CUDA_NAMESPACE {
reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
}
return cudaGetLastError();
return cuda_last_error();
}
template <bool Dequantize>
static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
@@ -535,7 +549,7 @@ namespace SD_CUDA_NAMESPACE {
dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
}
template <bool Swiglu, bool Dequantize, bool Rotated>
static cudaError_t quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, cudaStream_t stream) {
static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
const int64_t hidden = p.n / 2;
auto kernel = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
static std::mutex mutex;
@@ -543,17 +557,15 @@ namespace SD_CUDA_NAMESPACE {
{
std::lock_guard<std::mutex> lock(mutex);
if (!initialized[p.device.device]) {
#if !defined(SD_CUDA_VENDOR_HIP)
auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared_memory - 128));
if (error != cudaSuccess)
auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
if (error != cuda_success)
return error;
#endif
initialized[p.device.device] = true;
}
}
if (p.padded > p.rows) {
auto error = cudaMemsetAsync(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
if (error != cudaSuccess)
auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
if (error != cuda_success)
return error;
}
const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
@@ -562,7 +574,7 @@ namespace SD_CUDA_NAMESPACE {
kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
return cudaGetLastError();
return cuda_last_error();
}
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
const auto& p = *static_cast<const Plan*>(state);
@@ -570,7 +582,7 @@ 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_t>(context.stream);
auto stream = static_cast<CUDAStream>(context.stream);
auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
auto* dst = nodes[count - 1];
SDI8Operation op;
@@ -588,7 +600,7 @@ namespace SD_CUDA_NAMESPACE {
if (p.gemm && p.kind == SWIGLU_QUANTIZE)
alias = false;
void* output = alias ? scratch + p.output : dst->data;
cudaError_t error = cudaSuccess;
CUDAError error = cuda_success;
if (!mat) {
auto* quant = static_cast<int8_t*>(output);
auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
@@ -603,8 +615,8 @@ namespace SD_CUDA_NAMESPACE {
error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
else {
if (p.padded > p.rows)
error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
if (error == cudaSuccess)
error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
if (error == cuda_success)
quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
}
} else {
@@ -613,12 +625,12 @@ namespace SD_CUDA_NAMESPACE {
if (op.packed || p.rows == p.padded)
quant = data;
else {
error = cudaMemcpyAsync(quant, data, size_t(p.k) * p.rows, cudaMemcpyDeviceToDevice, stream);
if (error == cudaSuccess)
error = cudaMemsetAsync(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
if (error == cuda_success)
error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
}
}
if (error != cudaSuccess || cudaGetLastError() != cudaSuccess)
if (error != cuda_success || cuda_last_error() != cuda_success)
return GGML_STATUS_FAILED;
const auto* ws = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
@@ -643,11 +655,24 @@ namespace SD_CUDA_NAMESPACE {
return GGML_STATUS_FAILED;
#else
const int32_t alpha = 1, beta = 0;
#if defined(SD_CUDA_VENDOR_HIP)
#if HIP_VERSION >= 60500000
constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
#else
constexpr auto compute_type = HIPBLAS_R_32I;
#endif
const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
int(p.n), int(p.padded), int(p.k), &alpha,
op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
&beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
if (status != HIPBLAS_STATUS_SUCCESS)
#else
const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
int(p.n), int(p.padded), int(p.k), &alpha,
op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
&beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
if (status != CUBLAS_STATUS_SUCCESS)
#endif
return GGML_STATUS_FAILED;
#endif
}
@@ -665,7 +690,7 @@ namespace SD_CUDA_NAMESPACE {
swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
else
swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
error = cudaGetLastError();
error = cuda_last_error();
} else if (!p.gemm) {
const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
@@ -674,14 +699,14 @@ namespace SD_CUDA_NAMESPACE {
} else {
dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
}
error = cudaGetLastError();
error = cuda_last_error();
}
if (error == cudaSuccess && prefix)
error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream);
if (error == cuda_success && prefix)
error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
}
if (error == cudaSuccess && alias)
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
}
src/core/cuda/registration.cu
+1 -1
View File
@@ -1,4 +1,4 @@
#include "backend.h"
#include "runtime.h"
#ifdef SD_CUDA_CUTE_ATTENTION
void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
src/core/cuda/rope.cu
+20 -15
View File
@@ -9,7 +9,7 @@
namespace SD_CUDA_NAMESPACE {
static __device__ __forceinline__ float sd_warp_sum(float v) {
for (int mask = 16; mask > 0; mask >>= 1)
v = __fadd_rn(v, __shfl_xor_sync(0xffffffff, v, mask, 32));
v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
return v;
}
template <int D, bool interleaved, typename T>
@@ -54,7 +54,7 @@ namespace SD_CUDA_NAMESPACE {
const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
float rotated = value;
if constexpr (interleaved) {
const float other = __shfl_xor_sync(0xffffffff, value, 1, 32);
const float other = cuda_shuffle_xor(value, 1, 32);
if (col < 2 * rot_half) {
const float* m = pe + 4 * (token * rot_half + col / 2);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
@@ -136,6 +136,11 @@ namespace SD_CUDA_NAMESPACE {
return value != 0;
}
static float param_f32(const ggml_tensor* node, int index) {
float value;
std::memcpy(&value, node->op_params + index, sizeof(value));
return value;
}
static bool rope_is_node(const ggml_tensor* node) {
return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
@@ -174,7 +179,7 @@ namespace SD_CUDA_NAMESPACE {
}
template <typename T, bool Interleaved>
static cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, T* output, float scale) {
static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
const auto* x = nodes[0]->src[0];
const auto* pe = nodes[2]->src[1];
const int64_t rows = ggml_nrows(x);
@@ -183,7 +188,7 @@ namespace SD_CUDA_NAMESPACE {
rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data), \
static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows, \
x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), ggml_get_op_params_f32(nodes[0], 0), scale); \
x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
break
switch (x->ne[0]) {
SD_ROPE(32);
@@ -195,13 +200,13 @@ namespace SD_CUDA_NAMESPACE {
SD_ROPE(224);
SD_ROPE(256);
default:
return cudaErrorInvalidValue;
return cuda_invalid_value;
}
#undef SD_ROPE
return cudaGetLastError();
return cuda_last_error();
}
cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
if (fp16) {
return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
: rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
@@ -223,7 +228,7 @@ namespace SD_CUDA_NAMESPACE {
const auto* scale = ggml_graph_node(g, first + 3);
const auto* copy = ggml_graph_node(g, first + 4);
if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && ggml_get_op_params_f32(scale, 1) == 0.f &&
scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 5;
}
@@ -280,15 +285,15 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<cudaStream_t>(context.stream);
auto stream = static_cast<CUDAStream>(context.stream);
const auto* rope = nodes[count == 1 ? 0 : 2];
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
auto* dst = nodes[count - 1];
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
cudaError_t error;
CUDAError error;
if (count > 1) {
error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? ggml_get_op_params_f32(nodes[3], 0) : 1.f);
error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
const auto* pe = rope->src[1];
@@ -300,11 +305,11 @@ namespace SD_CUDA_NAMESPACE {
auto kernel = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
error = cudaGetLastError();
error = cuda_last_error();
}
if (error == cudaSuccess && alias)
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
void register_rope(ggml_backend_dev_t device) {
src/core/cuda/rope.h
+3 -2
View File
@@ -1,10 +1,11 @@
#ifndef __SD_CORE_CUDA_ROPE_H__
#define __SD_CORE_CUDA_ROPE_H__
#include "backend.h"
#include "runtime.h"
namespace SD_CUDA_NAMESPACE {
bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
cudaError_t rms_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
}
#endif
src/core/cuda/runtime.h
+111
View File
@@ -0,0 +1,111 @@
#ifndef __SD_CORE_CUDA_RUNTIME_H__
#define __SD_CORE_CUDA_RUNTIME_H__
#include "ggml-ext.h"
#if defined(SD_CUDA_VENDOR_HIP)
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#define SD_CUDA_NAMESPACE sd_hip
#define SD_CUDA_NAME "hip"
#elif defined(SD_CUDA_VENDOR_MUSA)
#include <musa_fp16.h>
#include <musa_runtime.h>
#define SD_CUDA_NAMESPACE sd_musa
#define SD_CUDA_NAME "musa"
#else
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#define SD_CUDA_NAMESPACE sd_cuda
#define SD_CUDA_NAME "cuda"
#endif
namespace SD_CUDA_NAMESPACE {
#if defined(SD_CUDA_VENDOR_HIP)
using CUDAError = hipError_t;
using CUDAStream = hipStream_t;
static constexpr auto cuda_success = hipSuccess;
static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
static constexpr auto cuda_kind = GGML_EXT_GPU_HIP;
static CUDAError cuda_last_error() {
return hipGetLastError();
}
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
return hipMemsetAsync(dst, value, bytes, stream);
}
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
}
template <typename Kernel>
static CUDAError cuda_set_shared_memory(Kernel, int) {
// HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
return cuda_success;
}
#elif defined(SD_CUDA_VENDOR_MUSA)
using CUDAError = musaError_t;
using CUDAStream = musaStream_t;
static constexpr auto cuda_success = musaSuccess;
static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA;
static CUDAError cuda_last_error() {
return musaGetLastError();
}
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
return musaMemsetAsync(dst, value, bytes, stream);
}
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
}
template <typename Kernel>
static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}
#else
using CUDAError = cudaError_t;
using CUDAStream = cudaStream_t;
static constexpr auto cuda_success = cudaSuccess;
static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA;
static CUDAError cuda_last_error() {
return cudaGetLastError();
}
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
return cudaMemsetAsync(dst, value, bytes, stream);
}
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
}
template <typename Kernel>
static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}
#endif
template <typename T>
static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
#if defined(SD_CUDA_VENDOR_HIP)
return __shfl(value, lane, width);
#else
return __shfl_sync(0xffffffff, value, lane, width);
#endif
}
template <typename T>
static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
#if defined(SD_CUDA_VENDOR_HIP)
return __shfl_xor(value, lane, width);
#else
return __shfl_xor_sync(0xffffffff, value, lane, width);
#endif
}
static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
return query && query(backend, context) && context->kind == cuda_kind;
}
static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
ggml_ext_gpu_device_info info = {};
auto query = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
if (!query || !query(device, &info) || info.kind != cuda_kind)
return {cuda_kind, -1, 0, 0, false};
return info;
}
}
#endif
src/core/cuda/sage_attention.cu
+4 -3
View File
@@ -1,5 +1,5 @@
#include "backend.h"
#include "core/ggml_extensions.h"
#include "runtime.h"
#if __CUDACC_VER_MAJOR__ >= 12
@@ -353,8 +353,9 @@ namespace sd_cuda {
if (info.device < 0)
return;
int previous;
cudaDeviceProp properties;
cudaFuncAttributes attributes;
if (cudaGetDevice(&previous) != cudaSuccess ||
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
cudaSetDevice(info.device) != cudaSuccess)
return;
const auto error = cudaFuncGetAttributes(&attributes, sage_output);
@@ -363,7 +364,7 @@ namespace sd_cuda {
// PTX compiled for an older target retains that target's instruction guards after JIT.
static std::map<int, SageDevice> devices;
auto& data = devices[info.device];
data = {info.cc / 10, std::min(attributes.binaryVersion, attributes.ptxVersion)};
data = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)};
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.sage_attn.cuda.v1";
src/core/cuda/sol_attention.cu
+3 -2
View File
@@ -1,7 +1,7 @@
// Sol-Attn diagonal-threshold algorithm:
// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
#include "backend.h"
#include "core/ggml_extensions.h"
#include "runtime.h"
#if __CUDACC_VER_MAJOR__ >= 12
@@ -379,8 +379,9 @@ namespace sd_cuda {
if (info.device < 0)
return;
int previous;
cudaDeviceProp properties;
cudaFuncAttributes attributes;
if (info.cc < 800 ||
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
return;
const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
src/core/ggml_extensions.cpp
+2 -2
View File
@@ -2,7 +2,6 @@
#ifndef SD_USE_UPSTREAM_GGML
#include "core/cpu/registration.h"
#include "ggml-cpu.h"
#include "ggml-ext.h"
#include <algorithm>
@@ -335,7 +334,8 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
return;
}
auto device = ggml_backend_get_device(backend);
const bool cpu = ggml_backend_is_cpu(backend);
auto reg = ggml_backend_dev_backend_reg(device);
const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
static std::mutex mutex;
static std::unordered_set<ggml_backend_dev_t> registered;
std::lock_guard<std::mutex> lock(mutex);