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

refactor: register CuTe attention and fusion in sd.cpp

Browse Source
This commit is contained in:
leejet committed 2026-10-31 19:00:24 +08:00
1 parent c6a4666289
commit 4b274821e9
5 files changed
+946 -10
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/cute_attention.cu
+883
View File
@@ -0,0 +1,883 @@
/*
* Adapted from FlashAttention v2.8.3:
* https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
* Sources: kernel_traits.h, utils.h, softmax.h, flash_fwd_kernel.h.
* Copyright (c) 2023-2024, Tri Dao.
*
* BSD 3-Clause License
*
* Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* * Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* * Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* * Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <cutlass/numeric_conversion.h>
#include <climits>
#include <cstdlib>
#include <cstring>
#include <cute/tensor.hpp>
#include <mutex>
#include <type_traits>
#include <unordered_map>
#include "ggml-ext.h"
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));
return v;
}
static __device__ __forceinline__ float sd_warp_max(float v) {
for (int mask = 16; mask > 0; mask >>= 1)
v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
return v;
}
namespace sd_fattn_cute {
using namespace cute;
template <bool A_in_regs = false, bool B_in_regs = false, typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename Tensor4, typename TiledMma, typename TiledCopyA, typename TiledCopyB, typename ThrCopyA, typename ThrCopyB>
__forceinline__ __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, Tensor4 const& tCsB, TiledMma tiled_mma, TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
if (!A_in_regs) {
cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
}
if (!B_in_regs) {
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
}
#pragma unroll
for (int i = 0; i < size<2>(tCrA); ++i) {
if (i < size<2>(tCrA) - 1) {
if (!A_in_regs) {
cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
}
if (!B_in_regs) {
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
}
}
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
}
}
template <typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename TiledMma, typename TiledCopy, typename ThrCopy>
__forceinline__ __device__ void gemm_rs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, ThrCopy smem_thr_copy_B) {
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
#pragma unroll
for (int i = 0; i < size<2>(tCrA); ++i) {
if (i < size<2>(tCrA) - 1) {
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
}
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
}
}
template <typename Layout>
__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
static_assert(decltype(size<0>(acc_layout))::value == 4);
static_assert(decltype(rank(acc_layout))::value == 3);
auto l = logical_divide(acc_layout, Shape<_2>{});
return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
};
template <typename MMA_traits, typename Layout>
__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
using X = Underscore;
static_assert(decltype(size<0>(acc_layout))::value == 4);
static_assert(decltype(rank(acc_layout))::value == 3);
constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
static_assert(mma_shape_K == 8 || mma_shape_K == 16);
if constexpr (mma_shape_K == 8) {
return acc_layout;
} else {
auto l = logical_divide(acc_layout, Shape<X, X, _2>{});
return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
}
};
template <typename TensorT>
__device__ __forceinline__ auto to_half(const TensorT& src) {
constexpr int count = decltype(size(src))::value;
auto dst = make_tensor<cutlass::half_t>(src.layout());
cutlass::NumericArrayConverter<cutlass::half_t, float, count> convert;
*reinterpret_cast<cutlass::Array<cutlass::half_t, count>*>(dst.data()) =
convert(*reinterpret_cast<const cutlass::Array<float, count>*>(src.data()));
return dst;
}
template <bool Full, bool Clear, bool FullCols = true, typename Copy, typename Src, typename Dst, typename Id>
__device__ __forceinline__ void copy_rows(Copy copier, const Src& src, Dst& dst, const Id& id, int rows, int cols = 0) {
#pragma unroll
for (int m = 0; m < size<1>(src); ++m) {
#pragma unroll
for (int k = 0; k < size<2>(src); ++k) {
if ((Full || get<0>(id(0, m, k)) < rows) && (FullCols || get<1>(id(0, m, k)) < cols)) {
cute::copy(copier, src(_, m, k), dst(_, m, k));
} else if constexpr (Clear) {
cute::clear(dst(_, m, k));
}
}
}
}
template <int Rows>
struct Softmax {
decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
template <bool First, typename Scores, typename Output>
__device__ __forceinline__ void update(Scores& acc_s, Output& acc_o) {
auto scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
#pragma unroll
for (int m = 0; m < Rows; ++m) {
const float old_max = First ? -INFINITY : maximum(m);
float next_max = First ? scores(m, 0) : fmaxf(old_max, scores(m, 0));
#pragma unroll
for (int n = 1; n < size<1>(scores); ++n)
next_max = fmaxf(next_max, scores(m, n));
next_max = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 2));
next_max = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 1));
maximum(m) = next_max;
float row_sum = 0.f;
if constexpr (!First) {
row_sum = sum(m);
if (old_max != next_max) {
const float factor = exp2f((old_max - next_max) * 1.4426950408889634f);
row_sum *= factor;
#pragma unroll
for (int n = 0; n < size<1>(output); ++n)
output(m, n) *= factor;
}
}
#pragma unroll
for (int n = 0; n < size<1>(scores); ++n) {
const float value = exp2f(scores(m, n) * 1.4426950408889634f - next_max * 1.4426950408889634f);
scores(m, n) = value;
row_sum += value;
}
// Defer the four-lane sum reduction until the final output normalization.
sum(m) = row_sum;
}
}
template <typename Output>
__device__ __forceinline__ void normalize(Output& acc_o) {
auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
#pragma unroll
for (int m = 0; m < Rows; ++m) {
float total = sum(m);
total += __shfl_xor_sync(0xffffffff, total, 2);
total += __shfl_xor_sync(0xffffffff, total, 1);
const float inverse = 1.f / total;
#pragma unroll
for (int n = 0; n < size<1>(output); ++n)
output(m, n) *= inverse;
}
}
};
template <int BM, int BN, int D = 128, int Warps = 4>
struct Traits {
using Element = cutlass::half_t;
static constexpr int Threads = Warps * 32;
using Mma = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<Int<Warps>, _1, _1>>,
Tile<Int<16 * Warps>, _16, _16>>;
static constexpr int SmemWidth = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
: 16;
static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
: 1;
using SmemAtom = decltype(composition(Swizzle<SwizzleBits, 3, 3>{},
Layout<Shape<_8, Int<SmemWidth>>, Stride<Int<SmemWidth>, _1>>{}));
using QLayout = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BM>, Int<D>>{}));
using KVLayout = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BN>, Int<D>>{}));
using VtLayout = decltype(composition(KVLayout{}, make_layout(Shape<Int<D>, Int<BN>>{}, GenRowMajor{})));
using CopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, Element>;
using CopyAtomT = Copy_Atom<SM75_U16x8_LDSM_T, Element>;
using GlobalLayout = Layout<Shape<Int<Threads / (SmemWidth / 8)>, Int<SmemWidth / 8>>, Stride<Int<SmemWidth / 8>, _1>>;
using GlobalCopy = decltype(make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, Element>{},
GlobalLayout{},
Layout<Shape<_1, _8>>{}));
using OutputCopy = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
GlobalLayout{},
Layout<Shape<_1, _8>>{}));
static constexpr int SharedBytes = (BM + 2 * BN) * D * sizeof(Element);
};
struct Params {
const cutlass::half_t *q, *k, *v;
cutlass::half_t* o;
int nq, nk, heads, kv_heads, dim;
};
template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
__global__ __launch_bounds__(Warps * 32) void kernel(Params p) {
using T = Traits<BM, BN, D, Warps>;
using Element = typename T::Element;
extern __shared__ char storage[];
const int tid = threadIdx.x, m_block = blockIdx.x, batch = blockIdx.y, head = blockIdx.z;
const int64_t q_offset = (int64_t(batch) * p.heads + head) * p.nq * p.dim;
const int64_t kv_offset = (int64_t(batch) * p.kv_heads + head / (p.heads / p.kv_heads)) * p.nk * p.dim;
auto mQ = make_tensor(make_gmem_ptr(p.q + q_offset), make_shape(p.nq, p.dim), make_stride(p.dim, _1{}));
auto mK = make_tensor(make_gmem_ptr(p.k + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
auto mV = make_tensor(make_gmem_ptr(p.v + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
auto gQ = local_tile(mQ, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
auto gK = local_tile(mK, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
auto gV = local_tile(mV, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
auto sQ = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
auto sK = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
auto sV = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
auto sVt = make_tensor(sV.data(), typename T::VtLayout{});
auto sVplain = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
typename T::GlobalCopy global_copy;
auto thread_copy = global_copy.get_thread_slice(tid);
auto tQgQ = thread_copy.partition_S(gQ);
auto tQsQ = thread_copy.partition_D(sQ);
auto tKgK = thread_copy.partition_S(gK);
auto tKsK = thread_copy.partition_D(sK);
auto tVgV = thread_copy.partition_S(gV);
auto tVsV = thread_copy.partition_D(sV);
auto tQcQ = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
auto tKVcKV = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
typename T::Mma mma;
auto thread_mma = mma.get_thread_slice(tid);
auto rQ = thread_mma.partition_fragment_A(sQ);
auto rK = thread_mma.partition_fragment_B(sK);
auto rV = thread_mma.partition_fragment_B(sVplain);
auto acc_o = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
auto copyQ = make_tiled_copy_A(typename T::CopyAtom{}, mma);
auto copyK = make_tiled_copy_B(typename T::CopyAtom{}, mma);
auto copyV = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
auto threadQ = copyQ.get_thread_slice(tid);
auto threadK = copyK.get_thread_slice(tid);
auto threadV = copyV.get_thread_slice(tid);
auto tSsQ = threadQ.partition_S(sQ);
auto tSsK = threadK.partition_S(sK);
auto tOsV = threadV.partition_S(sVt);
auto score_id = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
// Process the partial K/V tile first; subsequent tiles need no score masking.
int block = (p.nk - 1) / BN;
copy_rows<false, true, FullDim>(global_copy, tQgQ, tQsQ, tQcQ, p.nq - m_block * BM, p.dim);
auto last_k = tKgK(_, _, _, block);
copy_rows<false, true, FullDim>(global_copy, last_k, tKsK, tKVcKV, p.nk - block * BN, p.dim);
cute::cp_async_fence();
clear(acc_o);
Softmax<2 * size<1>(acc_o)> softmax;
auto step = [&](auto first_flag, int n) {
constexpr bool First = decltype(first_flag)::value;
auto acc_s = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
clear(acc_s);
cute::cp_async_wait<0>();
__syncthreads();
auto v_tile = tVgV(_, _, _, n);
copy_rows<!First, First, FullDim>(global_copy, v_tile, tVsV, tKVcKV, p.nk - n * BN, p.dim);
cute::cp_async_fence();
gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
if constexpr (First) {
#pragma unroll
for (int i = 0; i < size(acc_s); ++i) {
if (get<1>(score_id(i)) >= p.nk - n * BN)
acc_s(i) = -INFINITY;
}
}
cute::cp_async_wait<0>();
__syncthreads();
// Do not commit an empty async group: the next iteration waits for this K tile.
if (n > 0) {
auto next_k = tKgK(_, _, _, n - 1);
copy_rows<true, false, FullDim>(global_copy, next_k, tKsK, tKVcKV, 0, p.dim);
cute::cp_async_fence();
}
softmax.template update<First>(acc_s, acc_o);
auto rP = to_half(acc_s);
auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
};
step(std::true_type{}, block);
for (--block; block >= 0; --block)
step(std::false_type{}, block);
softmax.normalize(acc_o);
auto rO = to_half(acc_o);
auto sO = make_tensor(sQ.data(), typename T::QLayout{});
auto copyO = make_tiled_copy_C(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{}, mma);
auto threadO = copyO.get_thread_slice(tid);
auto tOrO = threadO.retile_S(rO);
auto tOsO = threadO.partition_D(sO);
cute::copy(copyO, tOrO, tOsO);
auto mO = make_tensor(make_gmem_ptr(p.o + int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim),
make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
typename T::OutputCopy output_copy;
auto global_threadO = output_copy.get_thread_slice(tid);
auto from_sO = global_threadO.partition_S(sO);
auto to_gO = global_threadO.partition_D(gO);
auto output_id = global_threadO.partition_D(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
__syncthreads();
auto output_regs = make_tensor<Element>(shape(to_gO));
cute::copy(output_copy, from_sO, output_regs);
copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
}
}
template <int D, bool interleaved, typename T>
static __global__ void rms_norm_mul_rope_apply_fused(
const float* x,
const float* weight,
const float* pe,
T* dst,
int64_t heads,
int64_t tokens,
int64_t rows,
int64_t nb1,
int64_t nb2,
int64_t nb3,
int rot_half,
float eps,
float output_scale) {
const int lane = threadIdx.x;
const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
if (row >= rows) {
return;
}
const int64_t head = row % heads;
const int64_t token = (row / heads) % tokens;
const int64_t batch = row / (heads * tokens);
const float* x_row = x + head * nb1 + token * nb2 + batch * nb3;
float v[D / 32], sums[8] = {};
#pragma unroll
for (int j = 0; j < D / 32; ++j) {
v[j] = x_row[j * 32 + lane];
sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
}
// Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
const float sum = __fadd_rn(
__fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
__fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
const float scale = rsqrtf(sum / float(D) + eps);
T* out = dst + ((batch * heads + head) * tokens + token) * D;
#pragma unroll
for (int j = 0; j < D / 32; ++j) {
const int col = j * 32 + lane;
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);
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]);
rotated = lane % 2 == 0
? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
}
} else {
if (col < 2 * rot_half) {
const int partner = col < rot_half ? col + rot_half : col - rot_half;
const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
const float* m = pe + 4 * (token * rot_half + col % rot_half);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = col < rot_half
? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
: __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
}
}
if constexpr (std::is_same<T, half>::value) {
out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
} else {
out[col] = rotated;
}
}
}
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_attention(const half* src, int8_t* dst, float* scales, int64_t k, int64_t rows, float input_scale, float input_bias) {
const int64_t row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= rows)
return;
extern __shared__ float values[];
__shared__ float maxima[32];
const half* src_row = src + row * k;
int8_t* dst_row = dst + row * k;
const int lane_id = tid % 32;
{
const int lane = lane_id;
constexpr int group_size = 256;
for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
float v[8];
#pragma unroll
for (int j = 0; j < 8; ++j) {
const float value = __fmaf_rn(__half2float(src_row[group * group_size + j * 32 + lane]), input_scale, input_bias);
v[j] = value * (1.0f / 16.0f);
}
// One logical warp owns H256, keeping every butterfly in registers.
#pragma unroll
for (int stride = 1; stride <= 4; stride *= 4) {
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);
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);
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
#pragma unroll
for (int t = 0; t < 4; ++t) {
values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
}
}
}
__syncthreads();
}
float local_max = 0.0f;
for (int64_t i = tid; i < k; i += blockDim.x) {
local_max = fmaxf(local_max, fabsf(values[i]));
}
local_max = sd_warp_max(local_max);
const int lane = tid & 31;
const int warp = tid >> 5;
if (lane == 0) {
maxima[warp] = local_max;
}
__syncthreads();
if (warp == 0) {
local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
local_max = sd_warp_max(local_max);
if (lane == 0) {
maxima[0] = local_max / 127.0f;
if (scales != nullptr) {
scales[row] = maxima[0];
} else {
*((float*)(dst_row + k)) = maxima[0];
}
}
}
__syncthreads();
const float row_scale = maxima[0];
const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
for (int64_t i = tid; i < k; i += blockDim.x) {
int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
value = max(-127, min(127, value));
dst_row[i] = (int8_t)value;
}
}
template <typename T>
static __global__ void copy_to_half(const char* src, half* dst, int64_t count, int64_t dim, int64_t rows, int64_t heads, size_t nb1, size_t nb2, size_t nb3, float scale) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i < count) {
const int64_t row = i / dim;
const auto* p = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
dst[i] = __float2half_rn(float(p[i % dim]) * scale);
}
}
static __global__ void copy_from_half(const half* src, float* dst, int64_t count, float scale, float bias, bool scaled) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
if (i < 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 = cuda_device(device);
const auto* q = node->src[0];
const auto* k = node->src[1];
const auto* v = node->src[2];
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->properties.sharedMemPerBlockOptin ||
q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
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) ||
(v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_F32) ||
k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type))
return false;
return k->ne[0] == q->ne[0] && q->ne[1] <= INT_MAX / q->ne[0] && k->ne[1] <= INT_MAX / q->ne[0] &&
q->ne[2] <= 65535 && q->ne[3] <= 65535 && ggml_nelements(q) <= INT_MAX && ggml_nelements(k) <= INT_MAX;
}
static bool attention_preferred(ggml_backend_dev_t, const ggml_tensor* node, void*) {
const auto* q = node->src[0];
const auto* k = node->src[1];
if (q->ne[0] == 40 || q->ne[0] == 72)
return q->ne[1] >= 512 && k->ne[1] >= 512;
if (q->ne[0] == 80 || q->ne[0] == 96 || q->ne[0] == 112)
return false;
const int minimum = q->ne[0] == 128 ? 2048 : 4096;
return q->ne[1] >= minimum && k->ne[1] >= minimum;
}
static bool attention_selected(ggml_backend_dev_t device, const ggml_tensor* node) {
return node->op == GGML_OP_FLASH_ATTN_EXT && attention_supports(device, node, nullptr) &&
(attention_preferred(device, node, nullptr) || !ggml_ext_dev_supports_native_op(device, node));
}
static bool fusion_enabled() {
static const bool disabled = [] {
const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
return value && std::atoi(value) != 0;
}();
return !disabled;
}
static bool query_fusion(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || rope->op != GGML_OP_ROPE_APPLY ||
mul->src[0] != norm || rope->src[0] != mul)
return false;
const auto* x = norm->src[0];
const auto* weight = mul->src[1];
const auto* pe = rope->src[1];
return x && weight && pe && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && rope->type == GGML_TYPE_F32 &&
ggml_are_same_shape(norm, mul) && weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] > 0 && pe->ne[2] <= x->ne[0] / 2 && pe->ne[3] == x->ne[2] &&
x->nb[0] == sizeof(float) && ggml_is_contiguous(weight) && ggml_is_contiguous(pe) && ggml_is_contiguous(rope);
}
static const ggml_tensor* graph_node(const ggml_cgraph* graph, int index) {
return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
}
static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, int first) {
const auto* attn = graph_node(graph, first);
const auto* previous = attn;
const int total = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
bool scaled = false;
for (int count = 1; count < 8 && first + count < total; ++count) {
const auto* node = graph_node(graph, first + count);
if (node->src[0] != previous)
break;
if (node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE) {
if (node->view_offs || !ggml_is_contiguous(node) || ggml_nelements(node) != ggml_nelements(attn))
break;
} else if (node->op == GGML_OP_SCALE && !scaled) {
if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
break;
scaled = true;
} else if (node->op == GGML_OP_QUANTIZE_I8_CONVROT) {
if (node->op_params[0] == 256 && previous->ne[0] % 256 == 0 &&
size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
return count + 1;
break;
} else
break;
previous = node;
}
if (first + 1 < total) {
const auto* next = graph_node(graph, first + 1);
if (next->op == GGML_OP_SCALE && next->src[0] == attn && ggml_is_contiguous(next) && ggml_are_same_shape(attn, next))
return 2;
}
return 1;
}
static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
if (!fusion_enabled())
return 0;
int offset = 0;
if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
!query_fusion(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
return 0;
offset = 3;
}
if (!attention_selected(device, graph_node(graph, first + offset)))
return 0;
const int count = offset + output_fusion(device, graph, first + offset);
return count > 1 ? count : 0;
}
static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
if (!fusion_enabled())
return;
for (int i = 3; i < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)); ++i) {
const auto* attn = graph_node(graph, i);
if (!attention_selected(device, attn))
continue;
for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
if (attn->src[0] == graph_node(graph, j + 2) &&
query_fusion(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
ggml_ext_graph_delay(graph, j, 3, i);
break;
}
}
}
}
struct AttentionPlan {
int index;
size_t offsets[4];
};
static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
auto* plan = new AttentionPlan{};
*state = plan;
plan->index = nodes[0]->op == GGML_OP_RMS_NORM ? 3 : 0;
GGML_ASSERT(plan->index < count);
const auto* attn = nodes[plan->index];
size_t size = 0;
for (int i = 0; i < 4; ++i) {
plan->offsets[i] = size;
const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half);
size += (bytes + 255) / 256 * 256;
}
*workspace_size = size;
return GGML_STATUS_SUCCESS;
}
template <int D, bool FullDim>
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;
static std::unordered_map<int, bool> initialized;
{
std::lock_guard<std::mutex> lock(mutex);
if (!initialized[context.device]) {
auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
if (error != cudaSuccess)
return error;
initialized[context.device] = true;
}
}
kernel<<<dim3((p.nq + 63) / 64, batches, p.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(p);
return cudaGetLastError();
}
template <int D>
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);
}
template <bool Interleaved>
static void query_launch(cudaStream_t stream, ggml_tensor* const* nodes, half* output, float scale) {
const auto* x = nodes[0]->src[0];
const auto* pe = nodes[2]->src[1];
const int64_t rows = ggml_nrows(x);
#define SD_QUERY(D) \
case D: \
rms_norm_mul_rope_apply_fused<D, Interleaved, half><<<(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]), param_f32(nodes[0], 0), scale); \
break
switch (x->ne[0]) {
SD_QUERY(32);
SD_QUERY(64);
SD_QUERY(96);
SD_QUERY(128);
SD_QUERY(160);
SD_QUERY(192);
SD_QUERY(224);
SD_QUERY(256);
}
#undef SD_QUERY
}
static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
const auto& plan = *static_cast<AttentionPlan*>(state);
const auto* attn = nodes[plan.index];
auto device = ggml_backend_get_device(execution->backend);
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));
half* work[4];
for (int i = 0; i < 4; ++i)
work[i] = reinterpret_cast<half*>(base + plan.offsets[i]);
const half* inputs[3] = {work[0], work[1], work[2]};
for (int i = 0; i < 3; ++i) {
const auto* src = attn->src[i];
const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
if (i == 0 && plan.index == 3) {
if (nodes[2]->op_params[0])
query_launch<true>(stream, nodes, work[0], scale);
else
query_launch<false>(stream, nodes, work[0], scale);
} else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
inputs[i] = static_cast<const half*>(src->data);
} else {
auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
}
if (cudaGetLastError() != cudaSuccess)
return GGML_STATUS_FAILED;
}
const auto* q = attn->src[0];
const auto* k = attn->src[1];
sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
cudaError_t error = cudaErrorInvalidValue;
switch (padded_dim(p.dim)) {
case 64:
error = attention_launch<64>(context, p, q->ne[3]);
break;
case 80:
error = attention_launch<80, true>(context, p, q->ne[3]);
break;
case 96:
error = attention_launch<96, true>(context, p, q->ne[3]);
break;
case 112:
error = attention_launch<112, true>(context, p, q->ne[3]);
break;
case 128:
error = attention_launch<128>(context, p, q->ne[3]);
break;
case 192:
error = attention_launch<192>(context, p, q->ne[3]);
break;
case 256:
error = attention_launch<256>(context, p, q->ne[3]);
break;
}
if (error != cudaSuccess)
return GGML_STATUS_FAILED;
float scale = 1.0f, bias = 0.0f;
for (int i = plan.index + 1; i < count; ++i) {
if (nodes[i]->op == GGML_OP_SCALE) {
scale = param_f32(nodes[i], 0);
bias = param_f32(nodes[i], 1);
}
}
auto* dst = nodes[count - 1];
if (dst->op == GGML_OP_QUANTIZE_I8_CONVROT) {
const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4;
auto* packed = static_cast<int8_t*>(dst->data);
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(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);
} else {
copy_from_half<<<(ggml_nelements(attn) + 255) / 256, 256, 0, stream>>>(work[3], static_cast<float*>(dst->data), ggml_nelements(attn), scale, bias, dst->op == GGML_OP_SCALE);
}
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}
static void attention_destroy(void* state, void*) {
delete static_cast<AttentionPlan*>(state);
}
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
if (!cuda_device(device))
return;
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.cute_attention.cuda.v1";
desc.op = GGML_OP_FLASH_ATTN_EXT;
desc.priority = 200;
desc.supports = attention_supports;
desc.preferred = attention_preferred;
desc.match = attention_match;
desc.optimize = attention_optimize;
desc.prepare = attention_prepare;
desc.execute = attention_execute;
desc.destroy = attention_destroy;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
desc.name = "sd.rms_rope_cute_attention.cuda.v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 300;
desc.supports = nullptr;
desc.preferred = nullptr;
desc.optimize = nullptr;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
}
src/core/ggml_extensions.cpp
+15 -3
View File
@@ -8,6 +8,10 @@
#include <mutex>
#include <unordered_set>
#ifdef SD_CUDA_EXTENSIONS
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
#endif
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
if (count != 2) {
return false;
@@ -155,9 +159,11 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!backend) {
return;
}
auto device = ggml_backend_get_device(backend);
auto reg = ggml_backend_dev_backend_reg(device);
if (!ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
auto device = ggml_backend_get_device(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;
const bool cuda = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr;
if (!cpu && !cuda) {
return;
}
static std::mutex mutex;
@@ -166,6 +172,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!registered.insert(device).second) {
return;
}
if (cuda) {
#ifdef SD_CUDA_EXTENSIONS
sd_ggml_register_cuda_attention(device);
#endif
return;
}
ggml_ext_kernel_desc desc = {};
desc.abi_version = GGML_EXT_ABI_VERSION;
desc.name = "sd.rope_apply.cpu.v1";