mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 03:49:09 +00:00
388 lines
18 KiB
Plaintext
388 lines
18 KiB
Plaintext
#include "core/ggml_extensions.h"
|
|
#include "runtime.h"
|
|
|
|
#if __CUDACC_VER_MAJOR__ >= 12
|
|
|
|
#include <algorithm>
|
|
#include <climits>
|
|
#include <cstring>
|
|
#include <map>
|
|
#include "sageattention/attn.cuh"
|
|
#include "sageattention/attn_fp8.cuh"
|
|
|
|
static constexpr int sage_cuda_version = __CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10;
|
|
|
|
namespace sd_cuda {
|
|
|
|
struct SageDevice {
|
|
int cc;
|
|
int compiled_cc;
|
|
};
|
|
|
|
static ggml_sage_attn_mode sage_mode(const SageDevice& device, ggml_sage_attn_mode requested) {
|
|
if (requested != GGML_SAGE_ATTN_AUTO)
|
|
return requested;
|
|
if (sage_cuda_version >= 12040 && device.cc >= 89 && device.compiled_cc >= 89) {
|
|
if (sage_cuda_version >= 12080 && device.cc != 90)
|
|
return GGML_SAGE_ATTN_2_PLUS_PLUS;
|
|
return GGML_SAGE_ATTN_2;
|
|
}
|
|
return GGML_SAGE_ATTN_AUTO;
|
|
}
|
|
|
|
static __device__ __forceinline__ float warp_reduce_max(float x) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset >>= 1)
|
|
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset));
|
|
return x;
|
|
}
|
|
|
|
static __global__ void sage_output(const half* input, float* output, int count) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i < count)
|
|
output[i] = __half2float(input[i]);
|
|
}
|
|
|
|
struct SageWorkspace {
|
|
size_t q;
|
|
size_t k;
|
|
size_t q_scale;
|
|
size_t k_scale;
|
|
size_t k_mean;
|
|
size_t v;
|
|
size_t v_scale;
|
|
size_t output;
|
|
size_t size;
|
|
|
|
explicit SageWorkspace(const ggml_tensor* dst) {
|
|
const ggml_tensor* qt = dst->src[0];
|
|
const ggml_tensor* kt = dst->src[1];
|
|
size = 0;
|
|
q = reserve(ggml_nelements(qt));
|
|
k = reserve(ggml_nelements(kt));
|
|
q_scale = reserve(GGML_PAD(qt->ne[1], 128) / 32 * 8 * qt->ne[2] * qt->ne[3] * sizeof(float));
|
|
k_scale = reserve(GGML_PAD(kt->ne[1], 64) / 64 * 4 * kt->ne[2] * kt->ne[3] * sizeof(float));
|
|
k_mean = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
|
|
v = reserve(kt->ne[0] * GGML_PAD(kt->ne[1], 64) * kt->ne[2] * kt->ne[3]);
|
|
v_scale = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
|
|
output = reserve(ggml_nelements(dst) * sizeof(half));
|
|
}
|
|
|
|
size_t reserve(size_t bytes) {
|
|
const size_t offset = size;
|
|
size += GGML_PAD(bytes, 256);
|
|
return offset;
|
|
}
|
|
};
|
|
|
|
template <int D>
|
|
static __global__ void sage_key_mean(const float* k, float* mean, int length) {
|
|
const int channel = blockIdx.x * 32 + threadIdx.x % 32;
|
|
const int row = threadIdx.x / 32;
|
|
const size_t base = size_t(blockIdx.y) * length * D;
|
|
float sum = 0.0f;
|
|
for (int token = row; token < length; token += 8) {
|
|
sum += k[base + size_t(token) * D + channel];
|
|
}
|
|
__shared__ float sums[256];
|
|
sums[threadIdx.x] = sum;
|
|
__syncthreads();
|
|
if (row == 0) {
|
|
for (int i = 1; i < 8; ++i) {
|
|
sum += sums[threadIdx.x + i * 32];
|
|
}
|
|
mean[size_t(blockIdx.y) * D + channel] = sum / length;
|
|
}
|
|
}
|
|
|
|
template <int D, int TOKENS, bool SMOOTH>
|
|
static __global__ void sage_quantize(const float* input, const float* mean, int8_t* output, float* scales, int length) {
|
|
constexpr int VALUES = D * TOKENS / 256;
|
|
const int head = blockIdx.y;
|
|
const int first = blockIdx.x * TOKENS * D;
|
|
const size_t base = size_t(head) * length * D;
|
|
float values[VALUES];
|
|
float amax = 1e-7f;
|
|
#pragma unroll
|
|
for (int i = 0; i < VALUES; ++i) {
|
|
const int index = first + threadIdx.x + i * 256;
|
|
float value = 0.0f;
|
|
if (index < length * D) {
|
|
value = input[base + index];
|
|
if constexpr (SMOOTH) {
|
|
value -= mean[size_t(head) * D + index % D];
|
|
}
|
|
}
|
|
values[i] = value;
|
|
amax = fmaxf(amax, fabsf(value));
|
|
}
|
|
amax = warp_reduce_max(amax);
|
|
__shared__ float maxima[8];
|
|
if (threadIdx.x % 32 == 0) {
|
|
maxima[threadIdx.x / 32] = amax;
|
|
}
|
|
__syncthreads();
|
|
if (threadIdx.x < 32) {
|
|
amax = threadIdx.x < 8 ? maxima[threadIdx.x] : 1e-7f;
|
|
amax = warp_reduce_max(amax);
|
|
if (threadIdx.x == 0) {
|
|
maxima[0] = amax;
|
|
scales[size_t(head) * gridDim.x + blockIdx.x] = amax / 127.0f;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
const float inverse_scale = 127.0f / maxima[0];
|
|
#pragma unroll
|
|
for (int i = 0; i < VALUES; ++i) {
|
|
const int index = first + threadIdx.x + i * 256;
|
|
if (index < length * D) {
|
|
output[base + index] = int8_t(__float2int_rn(values[i] * inverse_scale));
|
|
}
|
|
}
|
|
}
|
|
|
|
template <int D, bool KEY>
|
|
static __global__ void sage_quantize_per_thread(const float* input, const float* mean, int8_t* output, float* scales, int length) {
|
|
constexpr int GROUPS = KEY ? 4 : 8;
|
|
constexpr int TOKENS = KEY ? 64 : 32;
|
|
constexpr int VALUES = D * TOKENS / GROUPS / 32;
|
|
const int group = threadIdx.x / 32;
|
|
const int lane = threadIdx.x % 32;
|
|
const size_t base = size_t(blockIdx.y) * length * D;
|
|
float values[VALUES];
|
|
float amax = 0.0f;
|
|
#pragma unroll
|
|
for (int i = 0; i < VALUES; ++i) {
|
|
const int index = lane + i * 32;
|
|
const int row = index / D;
|
|
const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
|
|
float value = token < length ? input[base + size_t(token) * D + index % D] : 0.0f;
|
|
if constexpr (KEY) {
|
|
if (token < length) {
|
|
value -= mean[size_t(blockIdx.y) * D + index % D];
|
|
}
|
|
}
|
|
values[i] = value;
|
|
amax = fmaxf(amax, fabsf(value));
|
|
}
|
|
const float scale = warp_reduce_max(amax) / 127.0f + 1e-7f;
|
|
if (lane == 0) {
|
|
scales[(size_t(blockIdx.y) * gridDim.x + blockIdx.x) * GROUPS + group] = scale;
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < VALUES; ++i) {
|
|
const int index = lane + i * 32;
|
|
const int row = index / D;
|
|
const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
|
|
if (token < length) {
|
|
const float value = values[i] / scale;
|
|
output[base + size_t(token) * D + index % D] = int8_t(__float2int_rz(value + copysignf(0.5f, value)));
|
|
}
|
|
}
|
|
}
|
|
|
|
template <int D>
|
|
static __global__ void sage_value_scale(const half* v, float* scales, int length, float scale_max) {
|
|
const int channel = blockIdx.x * 32 + threadIdx.x % 32;
|
|
const int row = threadIdx.x / 32;
|
|
const size_t base = size_t(blockIdx.y) * length * D;
|
|
float amax = 1e-7f;
|
|
for (int token = row; token < length; token += 8) {
|
|
amax = fmaxf(amax, fabsf(__half2float(v[base + size_t(token) * D + channel])));
|
|
}
|
|
__shared__ float maxima[256];
|
|
maxima[threadIdx.x] = amax;
|
|
__syncthreads();
|
|
if (row == 0) {
|
|
for (int i = 1; i < 8; ++i) {
|
|
amax = fmaxf(amax, maxima[threadIdx.x + i * 32]);
|
|
}
|
|
scales[size_t(blockIdx.y) * D + channel] = amax / scale_max;
|
|
}
|
|
}
|
|
|
|
template <int D>
|
|
static __global__ void sage_quantize_value(const half* v, const float* scales, int8_t* output, int length, int padded) {
|
|
__shared__ uint8_t tile[32][D + 4];
|
|
const int head = blockIdx.y;
|
|
const int first = blockIdx.x * 32;
|
|
const size_t base = size_t(head) * length * D;
|
|
for (int i = threadIdx.x; i < 32 * D; i += 256) {
|
|
const int token = first + i / D;
|
|
const int channel = i % D;
|
|
const float value = token < length ? __half2float(v[base + size_t(token) * D + channel]) / scales[size_t(head) * D + channel] : 0.0f;
|
|
// FP8 MMA consumes pairs of tokens in a permuted, transposed V tile.
|
|
const int row = i / D;
|
|
const int permuted = row / 16 * 16 + row % 16 / 8 * 2 + row % 8 / 2 * 4 + row % 2;
|
|
tile[permuted][channel] = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
|
|
}
|
|
__syncthreads();
|
|
for (int i = threadIdx.x; i < 32 * D; i += 256) {
|
|
output[(size_t(head) * D + i / 32) * padded + first + i % 32] = tile[i % 32][i / 32];
|
|
}
|
|
}
|
|
|
|
template <int D, bool PLUS_PLUS>
|
|
static void sage_attn_fp8(cudaStream_t stream, ggml_tensor* dst, char* data, const SageWorkspace& work, float scale) {
|
|
const ggml_tensor* q = dst->src[0];
|
|
const ggml_tensor* k = dst->src[1];
|
|
const ggml_tensor* v = dst->src[2];
|
|
const int lq = q->ne[1], lk = k->ne[1], padded = GGML_PAD(lk, 64);
|
|
const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
|
|
auto* qi = reinterpret_cast<int8_t*>(data + work.q);
|
|
auto* ki = reinterpret_cast<int8_t*>(data + work.k);
|
|
auto* vi = reinterpret_cast<int8_t*>(data + work.v);
|
|
auto* qs = reinterpret_cast<float*>(data + work.q_scale);
|
|
auto* ks = reinterpret_cast<float*>(data + work.k_scale);
|
|
auto* vs = reinterpret_cast<float*>(data + work.v_scale);
|
|
auto* km = reinterpret_cast<float*>(data + work.k_mean);
|
|
auto* out = reinterpret_cast<half*>(data + work.output);
|
|
sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
|
|
sage_quantize_per_thread<D, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
|
|
static_cast<const float*>(q->data), nullptr, qi, qs, lq);
|
|
sage_quantize_per_thread<D, true><<<dim3(padded / 64, hk * batch), 128, 0, stream>>>(
|
|
static_cast<const float*>(k->data), km, ki, ks, lk);
|
|
// The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++.
|
|
sage_value_scale<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f);
|
|
sage_quantize_value<D><<<dim3(padded / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, vi, lk, padded);
|
|
auto kernel = sageattention::qk_int_sv_f8_attn_kernel<128, 64, 32, 64, D,
|
|
sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerThread,
|
|
sageattention::QuantGranularity::kPerThread, float, true, half,
|
|
sageattention::ComputeUnit::kCudaCore, sageattention::MaskMode::kNone, false, true, false, PLUS_PLUS>;
|
|
kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), 256 * D, stream>>>(
|
|
qi, ki, vi, out, nullptr, qs, ks, vs, nullptr, lq, lk, hq / hk,
|
|
D * lq * hq, D, D * lq, D * lk * hk, D, D * lk,
|
|
D * padded * hk, D * padded, padded, D * lq * hq, D * hq, D, scale);
|
|
sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
|
|
}
|
|
|
|
template <int D>
|
|
static void sage_attn(cudaStream_t stream, ggml_tensor* dst, char* data, ggml_sage_attn_mode mode, float scale) {
|
|
const ggml_tensor* q = dst->src[0];
|
|
const ggml_tensor* k = dst->src[1];
|
|
const ggml_tensor* v = dst->src[2];
|
|
const int lq = q->ne[1], lk = k->ne[1];
|
|
const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
|
|
const SageWorkspace work(dst);
|
|
if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS) {
|
|
sage_attn_fp8<D, true>(stream, dst, data, work, scale);
|
|
return;
|
|
}
|
|
if (mode == GGML_SAGE_ATTN_2) {
|
|
sage_attn_fp8<D, false>(stream, dst, data, work, scale);
|
|
return;
|
|
}
|
|
auto* qi = reinterpret_cast<int8_t*>(data + work.q);
|
|
auto* ki = reinterpret_cast<int8_t*>(data + work.k);
|
|
auto* qs = reinterpret_cast<float*>(data + work.q_scale);
|
|
auto* ks = reinterpret_cast<float*>(data + work.k_scale);
|
|
auto* km = reinterpret_cast<float*>(data + work.k_mean);
|
|
auto* out = reinterpret_cast<half*>(data + work.output);
|
|
|
|
sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
|
|
sage_quantize<D, 32, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
|
|
static_cast<const float*>(q->data), nullptr, qi, qs, lq);
|
|
sage_quantize<D, 64, true><<<dim3(GGML_PAD(lk, 64) / 64, hk * batch), 256, 0, stream>>>(
|
|
static_cast<const float*>(k->data), km, ki, ks, lk);
|
|
|
|
constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
|
|
auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
|
|
sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
|
|
sageattention::QuantGranularity::kPerWarp, float, false, half,
|
|
sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
|
|
kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
|
|
qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
|
|
lq, lk, hq / hk,
|
|
D * lq * hq, D, D * lq,
|
|
D * lk * hk, D, D * lk,
|
|
D * lk * hk, D, D * lk,
|
|
D * lq * hq, D * hq, D,
|
|
scale);
|
|
sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
|
|
}
|
|
|
|
static bool sage_supports(ggml_backend_dev_t, const ggml_tensor* op, void* user_data) {
|
|
const auto& device = *static_cast<const SageDevice*>(user_data);
|
|
SDSageAttentionParams params;
|
|
if (!sd_ggml_sage_attn_describe(op, ¶ms) || device.cc < 80 || device.compiled_cc < 80)
|
|
return false;
|
|
const auto mode = sage_mode(device, params.mode);
|
|
if (mode != GGML_SAGE_ATTN_AUTO && (sage_cuda_version < 12040 || device.cc < 89 || device.compiled_cc < 89))
|
|
return false;
|
|
if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS && sage_cuda_version < 12080)
|
|
return false;
|
|
const auto* q = op->src[0];
|
|
const auto* k = op->src[1];
|
|
const auto* v = op->src[2];
|
|
if (q->ne[0] != 64 && q->ne[0] != 128)
|
|
return false;
|
|
for (const auto* t : {q, k, v, op}) {
|
|
if (ggml_nelements(t) > INT_MAX || t->ne[2] * t->ne[3] > 65535 ||
|
|
(reinterpret_cast<uintptr_t>(t->data) & 15))
|
|
return false;
|
|
}
|
|
return k->ne[0] * GGML_PAD(k->ne[1], 64) * k->ne[2] * k->ne[3] <= INT_MAX;
|
|
}
|
|
|
|
static ggml_status sage_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
|
|
*workspace = SageWorkspace(nodes[0]).size;
|
|
return GGML_STATUS_SUCCESS;
|
|
}
|
|
|
|
static void sage_destroy(void*, void*) {}
|
|
|
|
static ggml_status sage_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void* user_data) {
|
|
ggml_ext_gpu_context context;
|
|
if (!cuda_context(execution->backend, &context))
|
|
return GGML_STATUS_FAILED;
|
|
SDSageAttentionParams params;
|
|
if (!sd_ggml_sage_attn_describe(nodes[0], ¶ms))
|
|
return GGML_STATUS_FAILED;
|
|
const auto mode = sage_mode(*static_cast<const SageDevice*>(user_data), params.mode);
|
|
auto* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
|
auto stream = static_cast<cudaStream_t>(context.stream);
|
|
if (nodes[0]->src[0]->ne[0] == 64)
|
|
sage_attn<64>(stream, nodes[0], data, mode, params.scale);
|
|
else
|
|
sage_attn<128>(stream, nodes[0], data, mode, params.scale);
|
|
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
|
}
|
|
|
|
void register_sage_attention(ggml_backend_dev_t device) {
|
|
const auto info = cuda_device_info(device);
|
|
if (info.device < 0)
|
|
return;
|
|
int previous;
|
|
cudaDeviceProp properties;
|
|
cudaFuncAttributes attributes;
|
|
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
|
|
cudaSetDevice(info.device) != cudaSuccess)
|
|
return;
|
|
const auto error = cudaFuncGetAttributes(&attributes, sage_output);
|
|
if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess)
|
|
return;
|
|
// 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 = {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";
|
|
desc.op = GGML_OP_CUSTOM;
|
|
desc.custom_op = sd_ggml_sage_attn_op();
|
|
desc.priority = 100;
|
|
desc.supports = sage_supports;
|
|
desc.prepare = sage_prepare;
|
|
desc.destroy = sage_destroy;
|
|
desc.execute = sage_execute;
|
|
desc.user_data = &data;
|
|
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
|
}
|
|
|
|
}
|
|
#else
|
|
namespace sd_cuda {
|
|
void register_sage_attention(ggml_backend_dev_t) {}
|
|
}
|
|
#endif
|