mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +00:00
410 lines
19 KiB
Plaintext
410 lines
19 KiB
Plaintext
// Sol-Attn diagonal-threshold algorithm:
|
|
// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
|
|
#include "core/ggml_extensions.h"
|
|
#include "runtime.h"
|
|
|
|
#if __CUDACC_VER_MAJOR__ >= 12
|
|
|
|
#include <cuda_bf16.h>
|
|
#include <algorithm>
|
|
#include <climits>
|
|
|
|
namespace sd_cuda {
|
|
|
|
static __device__ __forceinline__ float warp_reduce_sum(float x) {
|
|
#pragma unroll
|
|
for (int offset = 16; offset > 0; offset >>= 1)
|
|
x += __shfl_xor_sync(0xffffffff, x, offset);
|
|
return x;
|
|
}
|
|
|
|
static constexpr int SOL_DIM = 128;
|
|
static constexpr int SOL_BLOCK = 64;
|
|
|
|
struct SolWorkspace {
|
|
size_t q, k, v, kc, vc, qm, stats, threshold, size;
|
|
int tokens, blocks;
|
|
|
|
explicit SolWorkspace(const ggml_tensor* dst) {
|
|
const auto* input = dst->src[0];
|
|
const size_t heads = input->ne[2] * input->ne[3];
|
|
tokens = GGML_PAD(input->ne[1], SOL_BLOCK);
|
|
blocks = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
|
|
size = 0;
|
|
q = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
|
k = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
|
v = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
|
kc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
|
|
vc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
|
|
qm = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
|
|
stats = reserve(heads * 2 * SOL_DIM * sizeof(float));
|
|
threshold = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
|
|
}
|
|
|
|
size_t reserve(size_t bytes) {
|
|
const size_t offset = size;
|
|
size += GGML_PAD(bytes, 256);
|
|
return offset;
|
|
}
|
|
};
|
|
|
|
static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
|
|
const int index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (index >= padded * SOL_DIM) {
|
|
return;
|
|
}
|
|
const int token = index / SOL_DIM;
|
|
const int channel = index % SOL_DIM;
|
|
const size_t src = size_t(blockIdx.y) * length * SOL_DIM + index;
|
|
const size_t dst = size_t(blockIdx.y) * padded * SOL_DIM + index;
|
|
qb[dst] = __float2bfloat16(token < length ? q[src] : 0.f);
|
|
kb[dst] = __float2bfloat16(token < length ? k[src] : 0.f);
|
|
vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
|
|
__float2bfloat16(token < length ? v[src] : 0.f);
|
|
}
|
|
|
|
static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
|
|
const int channel = threadIdx.x;
|
|
const int block = blockIdx.x;
|
|
const int count = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
|
|
float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
|
|
for (int i = 0; i < count; ++i) {
|
|
const int token = block * SOL_BLOCK + i;
|
|
const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
|
|
sum_q += __bfloat162float(q[index]);
|
|
sum_k += __bfloat162float(k[index]);
|
|
sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
|
|
}
|
|
kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
|
|
__float2bfloat16(count ? sum_k / count : 0.f);
|
|
vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
|
|
if (count) {
|
|
qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
|
|
}
|
|
}
|
|
|
|
static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
|
|
const int channel = threadIdx.x;
|
|
float sum = 0.f, sum_sq = 0.f;
|
|
for (int i = 0; i < blocks; ++i) {
|
|
const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
|
|
sum += value;
|
|
sum_sq += value * value;
|
|
}
|
|
const float mean = sum / blocks;
|
|
stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel] = mean;
|
|
stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
|
|
}
|
|
|
|
static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
|
|
const int channel = threadIdx.x;
|
|
const float q = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
|
|
float mean = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
|
|
float var = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
|
|
mean = warp_reduce_sum(mean);
|
|
var = warp_reduce_sum(var);
|
|
__shared__ float means[4], vars[4];
|
|
if (channel % 32 == 0) {
|
|
means[channel / 32] = mean;
|
|
vars[channel / 32] = var;
|
|
}
|
|
__syncthreads();
|
|
if (channel == 0) {
|
|
mean = means[0] + means[1] + means[2] + means[3];
|
|
var = vars[0] + vars[1] + vars[2] + vars[3];
|
|
threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
|
|
mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
|
|
}
|
|
}
|
|
|
|
struct SolQTile {
|
|
uint32_t x[4];
|
|
};
|
|
struct SolKTile {
|
|
uint32_t x[2];
|
|
};
|
|
struct SolTile {
|
|
float x[4] = {};
|
|
static __device__ __forceinline__ int get_i(int i) {
|
|
return (i / 2) * 8 + threadIdx.x / 4;
|
|
}
|
|
static __device__ __forceinline__ int get_j(int i) {
|
|
return (threadIdx.x % 4) * 2 + i % 2;
|
|
}
|
|
};
|
|
|
|
static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
|
|
const auto* data = reinterpret_cast<const uint32_t*>(input);
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i)
|
|
tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
|
|
}
|
|
static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
|
|
const auto* data = reinterpret_cast<const uint32_t*>(input);
|
|
#pragma unroll
|
|
for (int i = 0; i < 2; ++i)
|
|
tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
|
|
}
|
|
static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
|
|
#if __CUDA_ARCH__ >= 800
|
|
auto* dst = reinterpret_cast<uint32_t*>(output.x);
|
|
asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
|
|
: "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
|
|
: "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
|
|
#else
|
|
asm("trap;");
|
|
#endif
|
|
}
|
|
|
|
static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
|
|
#pragma unroll
|
|
for (int n = 0; n < 8; ++n) {
|
|
scores[n] = SolTile{};
|
|
#pragma unroll
|
|
for (int d = 0; d < 8; ++d) {
|
|
SolKTile kt;
|
|
sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
|
|
sol_mma(scores[n], q[d], kt);
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
scores[n].x[i] *= scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <bool PROXY>
|
|
static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
|
|
float maximum[2] = {-INFINITY, -INFINITY};
|
|
#pragma unroll
|
|
for (int n = 0; n < 8; ++n) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
const int column = n * 8 + SolTile::get_j(i);
|
|
const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
|
|
: (start + column < length);
|
|
scores[n].x[i] = valid ? scores[n].x[i] : -INFINITY;
|
|
maximum[i / 2] = fmaxf(maximum[i / 2], scores[n].x[i]);
|
|
}
|
|
}
|
|
float alpha[2], sums[2] = {0.f, 0.f};
|
|
#pragma unroll
|
|
for (int row = 0; row < 2; ++row) {
|
|
maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
|
|
maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
|
|
const float next = fmaxf(row_max[row], maximum[row]);
|
|
alpha[row] = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
|
|
row_max[row] = next;
|
|
}
|
|
#pragma unroll
|
|
for (int n = 0; n < 8; ++n) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
const int column = n * 8 + SolTile::get_j(i);
|
|
const float p = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
|
|
const int count = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
|
|
sums[i / 2] += p * count;
|
|
probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int row = 0; row < 2; ++row) {
|
|
sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
|
|
sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
|
|
row_sum[row] = row_sum[row] * alpha[row] + sums[row];
|
|
}
|
|
__syncwarp();
|
|
#pragma unroll
|
|
for (int n = 0; n < 16; ++n) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
output[n].x[i] *= alpha[i / 2];
|
|
}
|
|
#pragma unroll
|
|
for (int d = 0; d < 4; ++d) {
|
|
SolQTile pt;
|
|
SolKTile vt;
|
|
sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
|
|
sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
|
|
sol_mma(output[n], pt, vt);
|
|
}
|
|
}
|
|
__syncwarp();
|
|
}
|
|
|
|
static __global__ __launch_bounds__(128) void sol_forward(
|
|
const nv_bfloat16* q,
|
|
const nv_bfloat16* k,
|
|
const nv_bfloat16* v,
|
|
const nv_bfloat16* kc,
|
|
const nv_bfloat16* vc,
|
|
const float* thresholds,
|
|
float* dst,
|
|
int length,
|
|
int padded,
|
|
int padded_blocks,
|
|
int heads,
|
|
float scale) {
|
|
#if __CUDA_ARCH__ >= 800
|
|
const int warp = threadIdx.y;
|
|
const int lane = threadIdx.x;
|
|
const int tid = warp * 32 + lane;
|
|
const int blocks = padded / SOL_BLOCK;
|
|
const int q_start = blockIdx.x * SOL_BLOCK;
|
|
const int q_len = min(SOL_BLOCK, length - q_start);
|
|
const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
|
|
SolQTile qt[8];
|
|
#pragma unroll
|
|
for (int d = 0; d < 8; ++d) {
|
|
sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
|
|
}
|
|
SolTile output[16];
|
|
float row_max[2] = {-INFINITY, -INFINITY};
|
|
float row_sum[2] = {0.f, 0.f};
|
|
__shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
|
|
__shared__ float routes[4][SOL_BLOCK];
|
|
__shared__ int exact[SOL_BLOCK];
|
|
const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
|
|
for (int group = 0; group < blocks; group += SOL_BLOCK) {
|
|
SolTile scores[8];
|
|
sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
|
|
#pragma unroll
|
|
for (int n = 0; n < 8; ++n) {
|
|
#pragma unroll
|
|
for (int column = 0; column < 2; ++column) {
|
|
float sum = scores[n].x[column] + scores[n].x[column + 2];
|
|
sum += __shfl_xor_sync(0xffffffff, sum, 4);
|
|
sum += __shfl_xor_sync(0xffffffff, sum, 8);
|
|
sum += __shfl_xor_sync(0xffffffff, sum, 16);
|
|
if (lane < 4) {
|
|
routes[warp][n * 8 + 2 * lane + column] = sum;
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
if (tid < SOL_BLOCK) {
|
|
const int block = group + tid;
|
|
const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
|
|
exact[tid] = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
|
|
}
|
|
__syncthreads();
|
|
sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
|
|
vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
|
|
padded_blocks, group, length, exact);
|
|
for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
|
|
if (!exact[block]) {
|
|
continue;
|
|
}
|
|
const int start = (group + block) * SOL_BLOCK;
|
|
sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
|
|
sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
|
|
v + base + start, padded, start, length, nullptr);
|
|
}
|
|
__syncthreads();
|
|
}
|
|
const int head = blockIdx.y % heads;
|
|
const int batch = blockIdx.y / heads;
|
|
#pragma unroll
|
|
for (int n = 0; n < 16; ++n) {
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
const int token = q_start + warp * 16 + SolTile::get_i(i);
|
|
if (token < length) {
|
|
const int channel = n * 8 + SolTile::get_j(i);
|
|
dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
|
|
__bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
|
|
}
|
|
}
|
|
}
|
|
#else
|
|
asm("trap;");
|
|
#endif
|
|
}
|
|
|
|
static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
|
|
SDSolAttentionParams params;
|
|
if (!sd_ggml_sol_attn_describe(op, ¶ms))
|
|
return false;
|
|
const auto* q = op->src[0];
|
|
if (q->ne[2] * q->ne[3] > 65535 || GGML_PAD(q->ne[1], SOL_BLOCK) > INT_MAX / SOL_DIM)
|
|
return false;
|
|
for (const auto* tensor : {op->src[0], op->src[1], op->src[2]})
|
|
if (reinterpret_cast<uintptr_t>(tensor->data) & 3)
|
|
return false;
|
|
return !(reinterpret_cast<uintptr_t>(op->data) & 3);
|
|
}
|
|
|
|
static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
|
|
*workspace = SolWorkspace(nodes[0]).size;
|
|
return GGML_STATUS_SUCCESS;
|
|
}
|
|
static void sol_destroy(void*, void*) {}
|
|
|
|
static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
|
|
ggml_ext_gpu_context context;
|
|
if (!cuda_context(execution->backend, &context))
|
|
return GGML_STATUS_FAILED;
|
|
auto* dst = nodes[0];
|
|
SDSolAttentionParams params;
|
|
if (!sd_ggml_sol_attn_describe(dst, ¶ms))
|
|
return GGML_STATUS_FAILED;
|
|
const auto* q = dst->src[0];
|
|
const SolWorkspace work(dst);
|
|
char* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
|
auto* qb = reinterpret_cast<nv_bfloat16*>(data + work.q);
|
|
auto* kb = reinterpret_cast<nv_bfloat16*>(data + work.k);
|
|
auto* vb = reinterpret_cast<nv_bfloat16*>(data + work.v);
|
|
auto* kc = reinterpret_cast<nv_bfloat16*>(data + work.kc);
|
|
auto* vc = reinterpret_cast<nv_bfloat16*>(data + work.vc);
|
|
auto* qm = reinterpret_cast<float*>(data + work.qm);
|
|
auto* stats = reinterpret_cast<float*>(data + work.stats);
|
|
auto* threshold = reinterpret_cast<float*>(data + work.threshold);
|
|
const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
|
|
const int blocks = work.tokens / SOL_BLOCK;
|
|
const float scale = params.scale * 1.4426950408889634f;
|
|
const auto stream = static_cast<cudaStream_t>(context.stream);
|
|
sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
|
|
static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
|
|
qb, kb, vb, length, work.tokens);
|
|
sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
|
|
sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
|
|
sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
|
|
sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
|
|
static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
|
|
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
|
}
|
|
|
|
void register_sol_attention(ggml_backend_dev_t device) {
|
|
const auto info = cuda_device_info(device);
|
|
if (info.device < 0)
|
|
return;
|
|
int previous;
|
|
cudaDeviceProp properties;
|
|
cudaFuncAttributes attributes;
|
|
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
|
|
cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
|
|
return;
|
|
const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
|
|
if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
|
|
std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
|
|
return;
|
|
ggml_ext_kernel_desc desc = {};
|
|
desc.abi_version = GGML_EXT_ABI_VERSION;
|
|
desc.name = "sd.sol_attn.cuda.v1";
|
|
desc.op = GGML_OP_CUSTOM;
|
|
desc.custom_op = sd_ggml_sol_attn_op();
|
|
desc.priority = 100;
|
|
desc.supports = sol_supports;
|
|
desc.prepare = sol_prepare;
|
|
desc.execute = sol_execute;
|
|
desc.destroy = sol_destroy;
|
|
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
|
}
|
|
|
|
}
|
|
#else
|
|
namespace sd_cuda {
|
|
void register_sol_attention(ggml_backend_dev_t) {}
|
|
}
|
|
#endif
|