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

chore: fotmat code

Browse Source
This commit is contained in:
leejet committed 2026-11-01 13:40:36 +08:00
1 parent 6edeac1b20
commit 5531c8878a
24 files changed
+3052 -3379
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/compute_workspace.cpp
+2 -2
View File
@@ -187,8 +187,8 @@ namespace sd {
} }
copy->buffer = buffer; copy->buffer = buffer;
// Capability checks inspect alignment even before weights are resident. // Capability checks inspect alignment even before weights are resident.
copy->data = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external)); copy->data = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
copy->extra = nullptr; copy->extra = nullptr;
} }
} }
auto copy_graph = ggml_new_graph_custom(context, graph_size, false); auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
src/core/cuda/affine.cu
+9 -9
View File
@@ -129,15 +129,15 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context; ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context)) if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED; return GGML_STATUS_FAILED;
auto stream = static_cast<CUDAStream>(context.stream); auto stream = static_cast<CUDAStream>(context.stream);
const auto* node = nodes[0]; const auto* node = nodes[0];
const auto* a = node->src[0]; const auto* a = node->src[0];
const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1]; const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
auto* dst = nodes[count - 1]; auto* dst = nodes[count - 1];
const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr; const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
const bool gate = residual != nullptr; const bool gate = residual != nullptr;
const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) || const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
(residual && overlaps(dst, residual) && dst->data != residual->data); (residual && overlaps(dst, residual) && dst->data != residual->data);
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
const int64_t count4 = ggml_nelements(dst) / 4; const int64_t count4 = ggml_nelements(dst) / 4;
const unsigned blocks = unsigned((count4 + 255) / 256); const unsigned blocks = unsigned((count4 + 255) / 256);
src/core/cuda/cute_attention.cu
+5 -5
View File
@@ -209,8 +209,8 @@ namespace sd_fattn_cute {
using Element = cutlass::half_t; using Element = cutlass::half_t;
static constexpr int Threads = Warps * 32; static constexpr int Threads = Warps * 32;
using Mma = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>, using Mma = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<Int<Warps>, _1, _1>>, Layout<Shape<Int<Warps>, _1, _1>>,
Tile<Int<16 * Warps>, _16, _16>>; Tile<Int<16 * Warps>, _16, _16>>;
static constexpr int SmemWidth = D % 64 == 0 ? 64 : D % 32 == 0 ? 32 static constexpr int SmemWidth = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
: 16; : 16;
static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2 static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
@@ -449,7 +449,7 @@ static __global__ void copy_to_half(const char* src, half* dst, int64_t count, i
if (i < count) { if (i < count) {
const int64_t row = i / dim; const int64_t row = i / dim;
const auto* p = reinterpret_cast<const T*>(src + (row % rows) * nb1 + const auto* p = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3); ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
dst[i] = __float2half_rn(float(p[i % dim]) * scale); dst[i] = __float2half_rn(float(p[i % dim]) * scale);
} }
} }
@@ -784,9 +784,9 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
desc.preferred = nullptr; desc.preferred = nullptr;
desc.optimize = nullptr; desc.optimize = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.i8_rope_cute_attention.cuda.v1"; desc.name = "sd.i8_rope_cute_attention.cuda.v1";
desc.op = GGML_OP_CUSTOM; desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_i8_op(); desc.custom_op = sd_ggml_i8_op();
desc.priority = 400; desc.priority = 400;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
} }
src/core/cuda/i8_gemm.cu
+2 -2
View File
@@ -197,8 +197,8 @@ namespace sd_i8_cutlass {
using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel; using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>; using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
using Kernel = std::conditional_t<StreamK, using Kernel = std::conditional_t<StreamK,
cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>, cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>; cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
}; };
struct Plan { struct Plan {
src/core/cuda/int8.cu
+4 -4
View File
@@ -661,10 +661,10 @@ namespace SD_CUDA_NAMESPACE {
#else #else
constexpr auto compute_type = HIPBLAS_R_32I; constexpr auto compute_type = HIPBLAS_R_32I;
#endif #endif
const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N, 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, 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), 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); &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
if (status != HIPBLAS_STATUS_SUCCESS) if (status != HIPBLAS_STATUS_SUCCESS)
#else #else
const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N, const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
src/core/cuda/sage_attention.cu
+3 -3
View File
@@ -287,9 +287,9 @@ namespace sd_cuda {
constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D; constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D, auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp, sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
sageattention::QuantGranularity::kPerWarp, float, false, half, sageattention::QuantGranularity::kPerWarp, float, false, half,
sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>; sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>( 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, qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
lq, lk, hq / hk, lq, lk, hq / hk,
src/core/cuda/sageattention/attn.cuh
+529 -623
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/sageattention/attn_fp8.cuh
+561 -657
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/sageattention/attn_utils.cuh
+644 -829
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/sageattention/cp_async.cuh
+96 -96
View File
@@ -24,17 +24,17 @@
namespace sageattention { namespace sageattention {
namespace cp_async { namespace cp_async {
enum class SharedMemFillMode { enum class SharedMemFillMode {
kFillZero, // Fill zero to shared memory when predicate is false kFillZero, // Fill zero to shared memory when predicate is false
kNoFill // Do not fill zero to shared memory when predicate is false kNoFill // Do not fill zero to shared memory when predicate is false
}; };
enum class PrefetchMode { enum class PrefetchMode {
kNoPrefetch, // Do not fetch additional data from global memory to L2 kNoPrefetch, // Do not fetch additional data from global memory to L2
kPrefetch // Fetch additional data from global memory to L2 kPrefetch // Fetch additional data from global memory to L2
}; };
#if (__CUDACC_VER_MAJOR__ >= 11) #if (__CUDACC_VER_MAJOR__ >= 11)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
@@ -42,107 +42,107 @@ enum class PrefetchMode {
#endif #endif
#endif #endif
/*! /*!
* \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
* cp.async instructions to a group * cp.async instructions to a group
*/ */
__device__ __forceinline__ void commit_group() { __device__ __forceinline__ void commit_group() {
#ifdef CP_ASYNC_ENABLED #ifdef CP_ASYNC_ENABLED
asm volatile("cp.async.commit_group;\n" ::); asm volatile("cp.async.commit_group;\n" ::);
#endif #endif
} }
/*! /*!
* \brief Wrapper of PTX cp.async.wait_group instruction * \brief Wrapper of PTX cp.async.wait_group instruction
* \tparam n Wait till most recent n groups are committed * \tparam n Wait till most recent n groups are committed
*/ */
template <size_t n> template <size_t n>
__device__ __forceinline__ void wait_group() { __device__ __forceinline__ void wait_group() {
#ifdef CP_ASYNC_ENABLED #ifdef CP_ASYNC_ENABLED
asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
#endif #endif
} }
/*! /*!
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
* global memory to shared memory * global memory to shared memory
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2 * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
* \tparam T Data type * \tparam T Data type
* \param smem_ptr Pointer to shared memory * \param smem_ptr Pointer to shared memory
* \param gmem_ptr Pointer to global memory * \param gmem_ptr Pointer to global memory
*/ */
template <PrefetchMode prefetch_mode, typename T> template <PrefetchMode prefetch_mode, typename T>
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
#ifdef CP_ASYNC_ENABLED #ifdef CP_ASYNC_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr)); uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
"l"(gmem_ptr), "n"(16), "r"(16)); "l"(gmem_ptr), "n"(16), "r"(16));
} else { } else {
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
"l"(gmem_ptr), "n"(16), "r"(16)); "l"(gmem_ptr), "n"(16), "r"(16));
} }
#else #else
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr); *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
#endif #endif
} }
/*! /*!
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
* global memory to shared memory with predicate. * global memory to shared memory with predicate.
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2 * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
* \tparam fill_mode Whether to fill zero to shared memory when predicate is false * \tparam fill_mode Whether to fill zero to shared memory when predicate is false
* \tparam T Data type * \tparam T Data type
* \param smem_ptr Pointer to shared memory * \param smem_ptr Pointer to shared memory
* \param gmem_ptr Pointer to global memory * \param gmem_ptr Pointer to global memory
* \param predicate Predicate value * \param predicate Predicate value
* \note fill zero is slower than not fill zero * \note fill zero is slower than not fill zero
*/ */
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T> template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
#ifdef CP_ASYNC_ENABLED #ifdef CP_ASYNC_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr)); uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
if constexpr (fill_mode == SharedMemFillMode::kFillZero) { if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
int src_in_bytes = predicate ? 16 : 0; int src_in_bytes = predicate ? 16 : 0;
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
} else { } else {
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
"l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
} }
} else { } else {
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
asm volatile( asm volatile(
"{\n" "{\n"
" .reg .pred p;\n" " .reg .pred p;\n"
" setp.ne.b32 p, %0, 0;\n" " setp.ne.b32 p, %0, 0;\n"
" @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
"}\n" ::"r"((int)predicate), "}\n" ::"r"((int)predicate),
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
} else { } else {
asm volatile( asm volatile(
"{\n" "{\n"
" .reg .pred p;\n" " .reg .pred p;\n"
" setp.ne.b32 p, %0, 0;\n" " setp.ne.b32 p, %0, 0;\n"
" @p cp.async.cg.shared.global [%1], [%2], %3;\n" " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
"}\n" ::"r"((int)predicate), "}\n" ::"r"((int)predicate),
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
} }
} }
#else #else
if (predicate) { if (predicate) {
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr); *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
} else { } else {
if constexpr (fill_mode == SharedMemFillMode::kFillZero) { if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
*((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0); *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
} }
} }
#endif #endif
} }
} // namespace cp_async } // namespace cp_async
} // namespace sageattention } // namespace sageattention
#endif #endif
src/core/cuda/sageattention/math.cuh
+118 -114
View File
@@ -26,137 +26,141 @@
namespace sageattention { namespace sageattention {
#ifndef USHORT_TYPE #ifndef USHORT_TYPE
#define USHORT_TYPE #define USHORT_TYPE
typedef unsigned short ushort; typedef unsigned short ushort;
#endif #endif
namespace math { namespace math {
// log2(e) // log2(e)
constexpr float log2e = 1.44269504088896340736f; constexpr float log2e = 1.44269504088896340736f;
constexpr float log2e_recp = 1.0f / log2e; constexpr float log2e_recp = 1.0f / log2e;
__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; } __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
return *(half2*)&x;
}
__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; } __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
return *(uint32_t*)&x;
}
/*! /*!
* \brief Wrapper of PTX ex2.approx instruction, which computes 2^x * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
* \param x input * \param x input
*/ */
__forceinline__ __device__ float ptx_exp2(float x) { __forceinline__ __device__ float ptx_exp2(float x) {
float y; float y;
asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
* \param x input * \param x input
*/ */
__forceinline__ __device__ float ptx_log2(float x) { __forceinline__ __device__ float ptx_log2(float x) {
float y; float y;
asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
* \param x input * \param x input
*/ */
__forceinline__ __device__ half2 ptx_exp2(half2 x) { __forceinline__ __device__ half2 ptx_exp2(half2 x) {
uint32_t y_u32; uint32_t y_u32;
uint32_t x_u32 = half2_as_uint32(x); uint32_t x_u32 = half2_as_uint32(x);
asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
return uint32_as_half2(y_u32); return uint32_as_half2(y_u32);
} }
/*! /*!
* \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
* \param x input * \param x input
*/ */
__forceinline__ __device__ half ptx_exp2(half x) { __forceinline__ __device__ half ptx_exp2(half x) {
ushort y_u16; ushort y_u16;
asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
return __ushort_as_half(y_u16); return __ushort_as_half(y_u16);
} }
/*! /*!
* \brief Wrapper of PTX rcp.approx instruction, which computes 1/x * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
* \param x input * \param x input
*/ */
__forceinline__ __device__ float ptx_rcp(float x) { __forceinline__ __device__ float ptx_rcp(float x) {
float y; float y;
asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
* between threads in a warp. * between threads in a warp.
* \param x The value in the source lane * \param x The value in the source lane
* \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta] * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
*/ */
__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) { __forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
float y; float y;
asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;" asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
: "=f"(y) : "=f"(y)
: "f"(x), "r"(lane_mask)); : "f"(x), "r"(lane_mask));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
* shuffle between threads in a warp. * shuffle between threads in a warp.
* \param x The value in the source lane * \param x The value in the source lane
* \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask] * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
*/ */
__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) { __forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
return __shfl_xor_sync(0xffffffff, x, lane_mask); return __shfl_xor_sync(0xffffffff, x, lane_mask);
} }
/*! /*!
* \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x) * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
* \param x input * \param x input
*/ */
__forceinline__ __device__ float rsqrt(float x) { __forceinline__ __device__ float rsqrt(float x) {
float y; float y;
asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x) * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
* \param x input * \param x input
*/ */
__forceinline__ __device__ float tanh(float x) { __forceinline__ __device__ float tanh(float x) {
float y; float y;
asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
return y; return y;
} }
/*! /*!
* \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x) * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
* \param x input * \param x input
*/ */
__forceinline__ __device__ half2 tanh(half2 x) { __forceinline__ __device__ half2 tanh(half2 x) {
uint32_t y_u32; uint32_t y_u32;
uint32_t x_u32 = half2_as_uint32(x); uint32_t x_u32 = half2_as_uint32(x);
asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
return uint32_as_half2(y_u32); return uint32_as_half2(y_u32);
} }
/*! /*!
* \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x) * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
* \param x input * \param x input
*/ */
__forceinline__ __device__ half tanh(half x) { __forceinline__ __device__ half tanh(half x) {
ushort y_u16; ushort y_u16;
asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
return __ushort_as_half(y_u16); return __ushort_as_half(y_u16);
} }
} // namespace math } // namespace math
} // namespace sageattention } // namespace sageattention
#endif #endif
src/core/cuda/sageattention/mma.cuh
+548 -594
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/sageattention/numeric_conversion.cuh
+80 -87
View File
@@ -38,118 +38,111 @@ namespace sageattention {
#define RUNTIME_ASSERT(x) assert(0 && x) #define RUNTIME_ASSERT(x) assert(0 && x)
#endif #endif
__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) { __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
uint16_t h0 = source & 0xFFFF; uint16_t h0 = source & 0xFFFF;
uint16_t h1 = (source >> 16) & 0xFFFF; uint16_t h1 = (source >> 16) & 0xFFFF;
asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0)); asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1)); asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
} }
__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t *dest, float *source0, float *source1) __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" \ "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" \ "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1])); : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t *dest, float *source0, float *source1) __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \ "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \ "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1])); : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \ "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \ "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \ "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \ "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo, hi;\n" \ ".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n" \ "mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
"cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo, hi;\n" \ ".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n" \ "mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
"cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ int8_t float_to_int8_rn(float x) __device__ __forceinline__ int8_t float_to_int8_rn(float x) {
{ uint32_t dst;
uint32_t dst; asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); return reinterpret_cast<const int8_t&>(dst);
return reinterpret_cast<const int8_t&>(dst); }
}
} // namespace sageattention } // namespace sageattention
#endif #endif
src/core/cuda/sageattention/permuted_smem.cuh
+154 -155
View File
@@ -23,180 +23,179 @@
#include <cuda_fp16.h> #include <cuda_fp16.h>
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include "cp_async.cuh" #include "cp_async.cuh"
#include "mma.cuh" #include "mma.cuh"
namespace sageattention { namespace sageattention {
enum class SwizzleMode { enum class SwizzleMode {
k32B, // for k32B mode, a line of shared memory must have 32B (16 half value) k32B, // for k32B mode, a line of shared memory must have 32B (16 half value)
k64B, // for k64B mode, a line of shared memory must have 64B (32 half value) k64B, // for k64B mode, a line of shared memory must have 64B (32 half value)
k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B. k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
}; };
// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
using b128_t = uint4; using b128_t = uint4;
/*! /*!
* \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
* that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
* permutation in the granularity of 128bit. * permutation in the granularity of 128bit.
* *
* This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless, * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
* meaning it does not maintain any information about the current pointer position. The offset returnd by * meaning it does not maintain any information about the current pointer position. The offset returnd by
* the struct can be used to access the shared memory through the provided interface. * the struct can be used to access the shared memory through the provided interface.
* *
* The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j). * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).
* We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp. * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
* Under this assumption, we do not permute for k32B mode. * Under this assumption, we do not permute for k32B mode.
*/ */
template <SwizzleMode swizzle_mode, uint32_t stride> template <SwizzleMode swizzle_mode, uint32_t stride>
struct smem_t { struct smem_t {
// The base pointer. // The base pointer.
b128_t* base; b128_t* base;
// How many b128_t value a row contains // How many b128_t value a row contains
// uint32_t stride; // uint32_t stride;
__device__ __forceinline__ smem_t() : base(nullptr) {} __device__ __forceinline__ smem_t() : base(nullptr) {}
template <typename T> template <typename T>
__device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) { __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
if constexpr (swizzle_mode == SwizzleMode::k128B) { if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode"); static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
} else if constexpr (swizzle_mode == SwizzleMode::k64B) { } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode"); static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
} else if constexpr (swizzle_mode == SwizzleMode::k32B) { } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode"); static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
} else { } else {
static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode"); static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
} }
} }
/*! /*!
* \brief Set the base pointer. * \brief Set the base pointer.
*/ */
template <typename T> template <typename T>
__device__ __forceinline__ void set_base(T* new_base) { __device__ __forceinline__ void set_base(T* new_base) {
base = (b128_t*)new_base; base = (b128_t*)new_base;
} }
/*! /*!
* \brief Compute the element offset given coordinates in a permuted shared memory. * \brief Compute the element offset given coordinates in a permuted shared memory.
* \param i The row index. * \param i The row index.
* \param j The column index. * \param j The column index.
*/ */
static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t &i, const uint32_t &j) { static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
if constexpr (swizzle_mode == SwizzleMode::k128B) { if constexpr (swizzle_mode == SwizzleMode::k128B) {
return i * stride + (j ^ (i % 8)); return i * stride + (j ^ (i % 8));
} else if constexpr (swizzle_mode == SwizzleMode::k64B) { } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
return i * stride + (j ^ ((i / 2) % 4)); return i * stride + (j ^ ((i / 2) % 4));
} else if constexpr (swizzle_mode == SwizzleMode::k32B) { } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
return i * stride + j; return i * stride + j;
} }
} }
/*! /*!
* \tparam step_size The step size to advance the offset in the permuted shared memory. * \tparam step_size The step size to advance the offset in the permuted shared memory.
* \param offset The current offset. * \param offset The current offset.
*/ */
template <uint32_t step_size> template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset) { static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
if constexpr (swizzle_mode == SwizzleMode::k128B) { if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size % 8 == 0, static_assert(step_size % 8 == 0,
"Unsupported step size"); "Unsupported step size");
return offset + step_size; return offset + step_size;
} else if constexpr (swizzle_mode == SwizzleMode::k64B) { } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
static_assert(step_size == 4, "Unsupported step size"); static_assert(step_size == 4, "Unsupported step size");
return offset + step_size; return offset + step_size;
} else if constexpr (swizzle_mode == SwizzleMode::k32B) { } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
static_assert(step_size == 2, "Unsupported step size"); static_assert(step_size == 2, "Unsupported step size");
return offset + step_size; return offset + step_size;
} }
} }
// ! use with care // ! use with care
template <uint32_t step_size> template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset, const uint32_t &step_idx) { static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
if constexpr (swizzle_mode == SwizzleMode::k128B) { if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0, static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
"Unsupported step size"); "Unsupported step size");
if constexpr (step_size == 2) { if constexpr (step_size == 2) {
return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8; return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
} else if constexpr (step_size == 4) { } else if constexpr (step_size == 4) {
return (offset ^ 0x4) + (step_idx % 2 == 1) * 8; return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
} else { } else {
// step_size % 8 == 0 // step_size % 8 == 0
return offset + step_size; return offset + step_size;
} }
} else if constexpr (swizzle_mode == SwizzleMode::k64B) { } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
static_assert(step_size == 2 || step_size == 4, "Unsupported step size"); static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
if constexpr (step_size == 2) { if constexpr (step_size == 2) {
return (offset ^ 0x2) + (step_idx % 2 == 1) * 4; return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
} else { } else {
return offset + step_size; return offset + step_size;
} }
} else if constexpr (swizzle_mode == SwizzleMode::k32B) { } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
return offset + step_size; return offset + step_size;
} }
} }
template <uint32_t step_size> template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t &offset) { static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
if constexpr (swizzle_mode == SwizzleMode::k128B) { if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
if constexpr (step_size == 4) { if constexpr (step_size == 4) {
return (offset ^ 0x4) + step_size * stride; return (offset ^ 0x4) + step_size * stride;
} else { } else {
// step_size % 8 == 0 // step_size % 8 == 0
return offset + step_size * stride; return offset + step_size * stride;
} }
} else if constexpr (swizzle_mode == SwizzleMode::k64B) { } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
if constexpr (step_size == 4) { if constexpr (step_size == 4) {
return (offset ^ 0x2) + step_size * stride; return (offset ^ 0x2) + step_size * stride;
} else { } else {
// step_size % 8 == 0 // step_size % 8 == 0
return offset + step_size * stride; return offset + step_size * stride;
} }
} else if constexpr (swizzle_mode == SwizzleMode::k32B) { } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
return offset + step_size * stride; return offset + step_size * stride;
} }
} }
__device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t &offset, uint32_t* R) const { __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset; b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x2(R, smem_ptr); mma::ldmatrix_m8n8x2(R, smem_ptr);
} }
__device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t &offset, uint32_t* R) const { __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset; b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x4(R, smem_ptr); mma::ldmatrix_m8n8x4(R, smem_ptr);
} }
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t &offset, uint32_t* R) const { __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset; b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x4_trans(R, smem_ptr); mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
} }
template <cp_async::SharedMemFillMode fill_mode, typename T> template <cp_async::SharedMemFillMode fill_mode, typename T>
__device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr, bool predicate) const { __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
b128_t* smem_ptr = base + offset; b128_t* smem_ptr = base + offset;
cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>( cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate); smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
} }
template <typename T> template <typename T>
__device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const { __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
b128_t* smem_ptr = base + offset; b128_t* smem_ptr = base + offset;
cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr)); cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
} }
template <typename T> template <typename T>
__device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const { __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
*reinterpret_cast<b128_t*>(gptr) = *(base + offset); *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
} }
}; };
} // namespace sageattention } // namespace sageattention
#endif #endif
src/core/ggml_extend.h
+1 -1
View File
@@ -4,9 +4,9 @@
#include <cstdint> #include <cstdint>
#include <vector> #include <vector>
#include "core/ggml_extensions.h"
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
#include "core/ggml_extensions.h"
#define EPS 1e-05f #define EPS 1e-05f
src/core/vulkan/int8.cpp
+1 -1
View File
@@ -613,7 +613,7 @@ namespace sd_vulkan_i8 {
auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst; auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst;
ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}}; ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}};
auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows, auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows,
std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) { std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
if (!dispatch_rows) if (!dispatch_rows)
return GGML_STATUS_SUCCESS; return GGML_STATUS_SUCCESS;
const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n; const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n;
src/core/vulkan/shaders/modulate_convrot.comp
+40 -29
View File
@@ -2,16 +2,28 @@
#extension GL_EXT_control_flow_attributes : require #extension GL_EXT_control_flow_attributes : require
layout(constant_id = 0) const uint K = 4096; layout(constant_id = 0) const uint K = 4096;
layout(constant_id = 1) const bool AFFINE = false; layout(constant_id = 1) const bool AFFINE = false;
layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer M { float data_m[]; }; float data_x[];
layout(binding = 2) readonly buffer H { float data_h[]; }; };
layout(binding = 3) writeonly buffer D { uint data_d[]; }; layout(binding = 1) readonly buffer M {
layout(binding = 4) writeonly buffer S { float data_s[]; }; float data_m[];
};
layout(binding = 2) readonly buffer H {
float data_h[];
};
layout(binding = 3) writeonly buffer D {
uint data_d[];
};
layout(binding = 4) writeonly buffer S {
float data_s[];
};
#ifdef RMS_AFFINE #ifdef RMS_AFFINE
layout(binding = 5) readonly buffer W { float data_w[]; }; layout(binding = 5) readonly buffer W {
float data_w[];
};
#endif #endif
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint rows; uint rows;
@@ -21,7 +33,8 @@ layout(push_constant) uniform parameter {
float eps; float eps;
uint k; uint k;
uint w_offset; uint w_offset;
} p; }
p;
shared float values[K]; shared float values[K];
shared vec2 sums[512]; shared vec2 sums[512];
@@ -36,7 +49,7 @@ float rms_row(uint row, uint tid) {
[[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) { [[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) {
float x = 0.0; float x = 0.0;
if (col < p.k) { if (col < p.k) {
x = data_x[p.x_offset + row * K + col]; x = data_x[p.x_offset + row * K + col];
values[col] = x; values[col] = x;
} }
sum += x * x; sum += x * x;
@@ -63,30 +76,31 @@ void main() {
for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) { for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
#ifdef RMS_AFFINE #ifdef RMS_AFFINE
float inv_std = rms_row(row, tid); float inv_std = rms_row(row, tid);
float mean = 0.0; float mean = 0.0;
#else #else
vec2 sum = vec2(0.0); vec2 sum = vec2(0.0);
[[unroll]] for (uint col = tid; col < K; col += 512) { [[unroll]] for (uint col = tid; col < K; col += 512) {
const float x = data_x[p.x_offset + row * K + col]; const float x = data_x[p.x_offset + row * K + col];
values[col] = x; values[col] = x;
if (!AFFINE) { if (!AFFINE) {
sum.x += x; sum.x += x;
sum.y = fma(x, x, sum.y); sum.y = fma(x, x, sum.y);
} }
} }
float mean = 0.0; float mean = 0.0;
float inv_std = 1.0; float inv_std = 1.0;
if (!AFFINE) { if (!AFFINE) {
sums[tid] = sum; sums[tid] = sum;
barrier(); barrier();
[[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
if (tid < stride) sums[tid] += sums[tid + stride]; if (tid < stride)
sums[tid] += sums[tid + stride];
barrier(); barrier();
} }
mean = sums[0].x / p.k; mean = sums[0].x / p.k;
// Match NORM's contraction even when precise propagates from the quantized output. // Match NORM's contraction even when precise propagates from the quantized output.
const float variance = fma(-mean, mean, sums[0].y / p.k); const float variance = fma(-mean, mean, sums[0].y / p.k);
inv_std = inversesqrt(variance + p.eps); inv_std = inversesqrt(variance + p.eps);
} }
#endif #endif
[[unroll]] for (uint col = tid; col < K; col += 512) { [[unroll]] for (uint col = tid; col < K; col += 512) {
@@ -96,7 +110,7 @@ void main() {
#endif #endif
if (AFFINE) { if (AFFINE) {
precise float modulated = values[col] * data_m[p.m_offset + col]; precise float modulated = values[col] * data_m[p.m_offset + col];
value = values[col] + modulated; value = values[col] + modulated;
value += data_h[p.h_offset + col]; value += data_h[p.h_offset + col];
} else { } else {
value = (values[col] - mean) * inv_std; value = (values[col] - mean) * inv_std;
@@ -107,15 +121,15 @@ void main() {
barrier(); barrier();
[[unroll]] for (uint stride = 1; stride < 256; stride *= 4) { [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) {
for (uint i = tid; i < K / 4; i += 512) { for (uint i = tid; i < K / 4; i += 512) {
const uint group = i / 64; const uint group = i / 64;
const uint lane = i % 64; const uint lane = i % 64;
const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride; const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride;
const float a = values[base]; const float a = values[base];
const float b = values[base + stride]; const float b = values[base + stride];
const float c = values[base + 2 * stride]; const float c = values[base + 2 * stride];
const float d = values[base + 3 * stride]; const float d = values[base + 3 * stride];
values[base] = a + b + c - d; values[base] = a + b + c - d;
values[base + stride] = a + b - c + d; values[base + stride] = a + b - c + d;
values[base + 2 * stride] = a - b + c + d; values[base + 2 * stride] = a - b + c + d;
values[base + 3 * stride] = -a + b + c + d; values[base + 3 * stride] = -a + b + c + d;
} }
@@ -133,16 +147,13 @@ void main() {
} }
barrier(); barrier();
} }
const float scale = sums[0].x / 127.0; const float scale = sums[0].x / 127.0;
const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale; const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale;
if (tid == 0) { if (tid == 0) {
data_s[row] = scale; data_s[row] = scale;
} }
for (uint i = tid; i < K / 4; i += 512) { for (uint i = tid; i < K / 4; i += 512) {
data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24);
| (quantize(values[i * 4 + 1], inv_scale) << 8)
| (quantize(values[i * 4 + 2], inv_scale) << 16)
| (quantize(values[i * 4 + 3], inv_scale) << 24);
} }
barrier(); barrier();
} }
src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp
+30 -18
View File
@@ -5,43 +5,55 @@
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
uint has_bias; uint has_bias;
uint row_offset; uint row_offset;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int data_a[]; }; int data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int data_a[];
layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; };
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
shared int weights[16 * 8]; shared int weights[16 * 8];
shared int activations[16 * 8]; shared int activations[16 * 8];
void main() { void main() {
const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x; const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x;
const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y; const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y;
const uint lx = gl_LocalInvocationID.x; const uint lx = gl_LocalInvocationID.x;
const uint ly = gl_LocalInvocationID.y; const uint ly = gl_LocalInvocationID.y;
const uint k_packed = p.k / 4; const uint k_packed = p.k / 4;
int sum = 0; int sum = 0;
for (uint k_base = 0; k_base < k_packed; k_base += 8) { for (uint k_base = 0; k_base < k_packed; k_base += 8) {
if (ly < 8) { if (ly < 8) {
weights[lx * 8 + ly] = output_index < p.n weights[lx * 8 + ly] = output_index < p.n
? data_w[output_index * k_packed + k_base + ly] ? data_w[output_index * k_packed + k_base + ly]
: 0; : 0;
} }
if (lx < 8) { if (lx < 8) {
activations[ly * 8 + lx] = row < p.rows activations[ly * 8 + lx] = row < p.rows
? data_a[row * k_packed + k_base + lx] ? data_a[row * k_packed + k_base + lx]
: 0; : 0;
} }
barrier(); barrier();
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp
+30 -18
View File
@@ -9,35 +9,47 @@
#extension GL_KHR_shader_subgroup_basic : require #extension GL_KHR_shader_subgroup_basic : require
layout(constant_id = 0) const uint BLOCK_SIZE = 32; layout(constant_id = 0) const uint BLOCK_SIZE = 32;
layout(constant_id = 1) const uint CM_M = 16; layout(constant_id = 1) const uint CM_M = 16;
layout(constant_id = 2) const uint CM_N = 16; layout(constant_id = 2) const uint CM_N = 16;
layout(constant_id = 3) const uint CM_K = 16; layout(constant_id = 3) const uint CM_K = 16;
const uint TILE_M = 4; const uint TILE_M = 4;
const uint TILE_N = 4; const uint TILE_N = 4;
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
uint has_bias; uint has_bias;
uint row_offset; uint row_offset;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; int8_t data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int8_t data_a[];
layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; };
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
shared int32_t accumulators[1024]; shared int32_t accumulators[1024];
void main() { void main() {
const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N; const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N;
const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M; const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M;
coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> sums[TILE_M * TILE_N]; coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> sums[TILE_M * TILE_N];
[[unroll]] [[unroll]]
for (uint i = 0; i < TILE_M * TILE_N; ++i) { for (uint i = 0; i < TILE_M * TILE_N; ++i) {
@@ -64,7 +76,7 @@ void main() {
[[unroll]] [[unroll]]
for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) { for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) {
const uint index = tile_m * TILE_N + tile_n; const uint index = tile_m * TILE_N + tile_n;
sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]); sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]);
} }
} }
} }
@@ -79,9 +91,9 @@ void main() {
gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) { for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) {
const uint row = row_base + tile_m * CM_M + index / CM_N; const uint row = row_base + tile_m * CM_M + index / CM_N;
const uint output_index = output_base + tile_n * CM_N + index % CM_N; const uint output_index = output_base + tile_n * CM_N + index % CM_N;
float value = float(accumulators[index]) * data_as[row] * data_ws[output_index]; float value = float(accumulators[index]) * data_as[row] * data_ws[output_index];
if (p.has_bias != 0) { if (p.has_bias != 0) {
value += data_b[output_index]; value += data_b[output_index];
} }
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp
+55 -36
View File
@@ -7,15 +7,14 @@
#extension GL_NV_cooperative_matrix2 : require #extension GL_NV_cooperative_matrix2 : require
#extension GL_KHR_memory_scope_semantics : require #extension GL_KHR_memory_scope_semantics : require
layout(constant_id = 0) const uint BN = 128; layout(constant_id = 0) const uint BN = 128;
layout(constant_id = 1) const uint ALIGNMENT = 0; layout(constant_id = 1) const uint ALIGNMENT = 0;
layout(constant_id = 2) const bool SWIGLU = false; layout(constant_id = 2) const bool SWIGLU = false;
layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 3) const uint BM = 128; layout(constant_id = 3) const uint BM = 128;
layout(constant_id = 4) const uint BK = 64; layout(constant_id = 4) const uint BK = 64;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
@@ -23,25 +22,45 @@ layout(push_constant) uniform parameter
uint row_offset; uint row_offset;
float output_scale; float output_scale;
uint group_cols; uint group_cols;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; int8_t data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int8_t data_a[];
};
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
// Store FP32 bit patterns without changing the integer accumulator fragment layout. // Store FP32 bit patterns without changing the integer accumulator fragment layout.
layout(binding = 5, std430) writeonly buffer D { int32_t data_d[]; }; layout(binding = 5, std430) writeonly buffer D {
int32_t data_d[];
};
#ifdef SPLIT_SWIGLU #ifdef SPLIT_SWIGLU
layout(binding = 6, std430) readonly buffer UW { int8_t data_uw[]; }; layout(binding = 6, std430) readonly buffer UW {
layout(binding = 7, std430) readonly buffer US { float data_us[]; }; int8_t data_uw[];
};
layout(binding = 7, std430) readonly buffer US {
float data_us[];
};
#endif #endif
#ifdef RESIDUAL #ifdef RESIDUAL
layout(binding = 6, std430) readonly buffer G { float data_g[]; }; layout(binding = 6, std430) readonly buffer G {
layout(binding = 7, std430) readonly buffer R { float data_r[]; }; float data_g[];
};
layout(binding = 7, std430) readonly buffer R {
float data_r[];
};
#endif #endif
int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
const in uint row_base, const in uint col_base) {
const uint r = row_base + col; const uint r = row_base + col;
const uint c = col_base + row; const uint c = col_base + row;
if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) { if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) {
@@ -75,32 +94,32 @@ int32_t dequantize_swiglu(const in uint row, const in uint col, const in int32_t
} }
void main() { void main() {
const uint k_size = p.k & ~255u; const uint k_size = p.k & ~255u;
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size); a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size);
w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size); w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size);
d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n); d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n);
tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0);
// Keep a weight panel in L2 while visiting the activation tiles. // Keep a weight panel in L2 while visiting the activation tiles.
const uint col_tiles = gl_NumWorkGroups.x; const uint col_tiles = gl_NumWorkGroups.x;
const uint row_tiles = gl_NumWorkGroups.y; const uint row_tiles = gl_NumWorkGroups.y;
const uint group_cols = clamp(p.group_cols, 1u, col_tiles); const uint group_cols = clamp(p.group_cols, 1u, col_tiles);
const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x; const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x;
const uint group = linear / (group_cols * row_tiles); const uint group = linear / (group_cols * row_tiles);
const uint first_col = group * group_cols; const uint first_col = group * group_cols;
const uint cols = min(col_tiles - first_col, group_cols); const uint cols = min(col_tiles - first_col, group_cols);
const uint within_group = linear - group * group_cols * row_tiles; const uint within_group = linear - group * group_cols * row_tiles;
const uint row = p.row_offset + (within_group / cols) * BN; const uint row = p.row_offset + (within_group / cols) * BN;
const uint col = (first_col + within_group % cols) * BM; const uint col = (first_col + within_group % cols) * BM;
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0); coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum_up = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0); coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum_up = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
if (ALIGNMENT != 0) { if (ALIGNMENT != 0) {
tensorLayoutNV<2> a_full = createTensorLayoutNV(2); tensorLayoutNV<2> a_full = createTensorLayoutNV(2);
tensorLayoutNV<2> w_full = createTensorLayoutNV(2); tensorLayoutNV<2> w_full = createTensorLayoutNV(2);
a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size); a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size);
w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size); w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size);
for (uint k = 0; k < p.k; k += BK) { for (uint k = 0; k < p.k; k += BK) {
coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w; coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a; coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
@@ -145,7 +164,7 @@ void main() {
// The aligned dispatch contains only complete tiles and a 128-aligned output width. // The aligned dispatch contains only complete tiles and a 128-aligned output width.
if (ALIGNMENT == 1) { if (ALIGNMENT == 1) {
tensorLayoutNV<2> d_full = createTensorLayoutNV(2); tensorLayoutNV<2> d_full = createTensorLayoutNV(2);
d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u); d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u);
coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose); coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose);
} else { } else {
coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose); coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose);
src/core/vulkan/shaders/quantize_i8_convrot.comp
+30 -27
View File
@@ -4,45 +4,51 @@
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint k; uint k;
uint rows; uint rows;
} p; }
p;
layout(binding = 0, std430) readonly buffer A { float data_a[]; }; layout(binding = 0, std430) readonly buffer A {
layout(binding = 1, std430) writeonly buffer D { uint data_d[]; }; float data_a[];
layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; };
layout(binding = 1, std430) writeonly buffer D {
uint data_d[];
};
layout(binding = 2, std430) writeonly buffer S {
float data_s[];
};
shared float values[256]; shared float values[256];
shared float maxima[64]; shared float maxima[64];
void transform_group(const uint row, const uint group) { void transform_group(const uint row, const uint group) {
const uint tid = gl_LocalInvocationID.x; const uint tid = gl_LocalInvocationID.x;
const uint base_offset = row * p.k + group * 256; const uint base_offset = row * p.k + group * 256;
[[unroll]] [[unroll]]
for (uint i = 0; i < 4; ++i) { for (uint i = 0; i < 4; ++i) {
const uint index = tid + i * 64; const uint index = tid + i * 64;
values[index] = data_a[base_offset + index] * (1.0 / 16.0); values[index] = data_a[base_offset + index] * (1.0 / 16.0);
} }
barrier(); barrier();
[[unroll]] [[unroll]]
for (uint stride = 1; stride < 256; stride *= 4) { for (uint stride = 1; stride < 256; stride *= 4) {
const uint base = (tid / stride) * 4 * stride + tid % stride; const uint base = (tid / stride) * 4 * stride + tid % stride;
const uint i0 = base; const uint i0 = base;
const uint i1 = i0 + stride; const uint i1 = i0 + stride;
const uint i2 = i1 + stride; const uint i2 = i1 + stride;
const uint i3 = i2 + stride; const uint i3 = i2 + stride;
const float a = values[i0]; const float a = values[i0];
const float b = values[i1]; const float b = values[i1];
const float c = values[i2]; const float c = values[i2];
const float d = values[i3]; const float d = values[i3];
values[i0] = a + b + c - d; values[i0] = a + b + c - d;
values[i1] = a + b - c + d; values[i1] = a + b - c + d;
values[i2] = a - b + c + d; values[i2] = a - b + c + d;
values[i3] = -a + b + c + d; values[i3] = -a + b + c + d;
barrier(); barrier();
} }
} }
@@ -53,9 +59,9 @@ uint quantize_value(const float value, const float inv_scale) {
} }
void quantize_row(const uint row) { void quantize_row(const uint row) {
const uint tid = gl_LocalInvocationID.x; const uint tid = gl_LocalInvocationID.x;
const uint groups = p.k / 256; const uint groups = p.k / 256;
float local_max = 0.0; float local_max = 0.0;
for (uint group = 0; group < groups; ++group) { for (uint group = 0; group < groups; ++group) {
transform_group(row, group); transform_group(row, group);
@@ -85,11 +91,8 @@ void quantize_row(const uint row) {
const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
for (uint group = 0; group < groups; ++group) { for (uint group = 0; group < groups; ++group) {
transform_group(row, group); transform_group(row, group);
const uint index = tid * 4; const uint index = tid * 4;
const uint packed = quantize_value(values[index], inv_scale) const uint packed = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
| (quantize_value(values[index + 1], inv_scale) << 8)
| (quantize_value(values[index + 2], inv_scale) << 16)
| (quantize_value(values[index + 3], inv_scale) << 24);
data_d[(row * p.k + group * 256) / 4 + tid] = packed; data_d[(row * p.k + group * 256) / 4 + tid] = packed;
barrier(); barrier();
} }
src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp
+48 -25
View File
@@ -8,12 +8,24 @@
#extension GL_EXT_control_flow_attributes : require #extension GL_EXT_control_flow_attributes : require
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter { uint k; uint rows; } p; layout(push_constant) uniform parameter {
layout(binding = 0, std430) readonly buffer A { float data_a[]; }; uint k;
layout(binding = 1, std430) writeonly buffer D { int8_t data_d[]; }; uint rows;
layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; }
p;
layout(binding = 0, std430) readonly buffer A {
float data_a[];
};
layout(binding = 1, std430) writeonly buffer D {
int8_t data_d[];
};
layout(binding = 2, std430) writeonly buffer S {
float data_s[];
};
#if defined(SWIGLU) || defined(SIGMOID_MUL) #if defined(SWIGLU) || defined(SIGMOID_MUL)
layout(binding = 3, std430) readonly buffer B { float data_b[]; }; layout(binding = 3, std430) readonly buffer B {
float data_b[];
};
#endif #endif
shared float maxima[8]; shared float maxima[8];
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
@@ -26,7 +38,7 @@ layout(constant_id = 0) const uint K = 4096;
vec4 rotate(float a, float b, float c, float d) { vec4 rotate(float a, float b, float c, float d) {
precise vec4 r = vec4(a + b + c - d, a + b - c + d, precise vec4 r = vec4(a + b + c - d, a + b - c + d,
a - b + c + d, -a + b + c + d); a - b + c + d, -a + b + c + d);
return r; return r;
} }
@@ -37,12 +49,12 @@ void transform_group(uint row, uint group, out float v[8]) {
float value = data_a[base + lane + j * 32]; float value = data_a[base + lane + j * 32];
#ifdef SWIGLU #ifdef SWIGLU
precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32]; precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32];
value = activated; value = activated;
#endif #endif
#ifdef SIGMOID_MUL #ifdef SIGMOID_MUL
precise float gate = 1.0 / (1.0 + exp(-value)); precise float gate = 1.0 / (1.0 + exp(-value));
precise float activated = data_b[base + lane + j * 32] * gate; precise float activated = data_b[base + lane + j * 32] * gate;
value = activated; value = activated;
#endif #endif
v[j] = value * (1.0 / 16.0); v[j] = value * (1.0 / 16.0);
} }
@@ -51,18 +63,21 @@ void transform_group(uint row, uint group, out float v[8]) {
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride), vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride),
subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride)); subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride));
v[j] = r[(lane / stride) & 3u]; v[j] = r[(lane / stride) & 3u];
} }
} }
[[unroll]] for (uint j = 0; j < 8; j += 2) { [[unroll]] for (uint j = 0; j < 8; j += 2) {
vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16), vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16),
subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16)); subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16));
v[j] = r[lane / 16]; v[j] = r[lane / 16];
v[j + 1] = r[2 + lane / 16]; v[j + 1] = r[2 + lane / 16];
} }
[[unroll]] for (uint j = 0; j < 2; ++j) { [[unroll]] for (uint j = 0; j < 2; ++j) {
vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]); vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]);
v[j] = r.x; v[j + 2] = r.y; v[j + 4] = r.z; v[j + 6] = r.w; v[j] = r.x;
v[j + 2] = r.y;
v[j + 4] = r.z;
v[j + 6] = r.w;
} }
} }
@@ -72,26 +87,32 @@ void main() {
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
float values[((K + 2047) / 2048) * 8]; float values[((K + 2047) / 2048) * 8];
[[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
if (gl_SubgroupID + g * 8 >= K / 256) continue; if (gl_SubgroupID + g * 8 >= K / 256)
float v[8]; transform_group(row, gl_SubgroupID + g * 8, v); continue;
float v[8];
transform_group(row, gl_SubgroupID + g * 8, v);
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
values[g * 8 + j] = v[j]; values[g * 8 + j] = v[j];
local_max = max(local_max, abs(v[j])); local_max = max(local_max, abs(v[j]));
} }
} }
#else #else
for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
float v[8]; transform_group(row, group, v); float v[8];
[[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j])); transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j)
local_max = max(local_max, abs(v[j]));
} }
#endif #endif
local_max = subgroupMax(local_max); local_max = subgroupMax(local_max);
if (gl_SubgroupInvocationID == 0) maxima[gl_SubgroupID] = local_max; if (gl_SubgroupInvocationID == 0)
maxima[gl_SubgroupID] = local_max;
barrier(); barrier();
if (gl_LocalInvocationID.x == 0) { if (gl_LocalInvocationID.x == 0) {
float row_max = 0.0; float row_max = 0.0;
for (uint j = 0; j < 8; ++j) row_max = max(row_max, maxima[j]); for (uint j = 0; j < 8; ++j)
maxima[0] = row_max; row_max = max(row_max, maxima[j]);
maxima[0] = row_max;
data_s[row] = row_max / 127.0; data_s[row] = row_max / 127.0;
} }
barrier(); barrier();
@@ -99,14 +120,16 @@ void main() {
const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
[[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
if (gl_SubgroupID + g * 8 >= K / 256) continue; if (gl_SubgroupID + g * 8 >= K / 256)
continue;
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127)); data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127));
} }
} }
#else #else
for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
float v[8]; transform_group(row, group, v); float v[8];
transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127)); data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127));
} }
src/core/vulkan/shaders/rms_rope_apply.comp
+32 -21
View File
@@ -5,16 +5,26 @@
#extension GL_EXT_control_flow_attributes : require #extension GL_EXT_control_flow_attributes : require
layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer PE { float data_pe[]; }; float data_x[];
};
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
#ifdef CAST_F16 #ifdef CAST_F16
#extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
layout(binding = 2) writeonly buffer D { float16_t data_d[]; }; layout(binding = 2) writeonly buffer D {
float16_t data_d[];
};
#else #else
layout(binding = 2) writeonly buffer D { float data_d[]; }; layout(binding = 2) writeonly buffer D {
float data_d[];
};
#endif #endif
layout(binding = 3) readonly buffer W { float data_w[]; }; layout(binding = 3) readonly buffer W {
float data_w[];
};
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint half_dim; uint half_dim;
@@ -31,10 +41,11 @@ layout(push_constant) uniform parameter {
float eps; float eps;
uint w_offset; uint w_offset;
float output_scale; float output_scale;
} pc; }
pc;
float shuffle_channel(vec4 values, uint channel) { float shuffle_channel(vec4 values, uint channel) {
const uint lane = channel & 31u; const uint lane = channel & 31u;
const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane), const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane),
subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane)); subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane));
return shuffled[channel / 32]; return shuffled[channel / 32];
@@ -45,43 +56,43 @@ void main() {
if (head >= pc.heads) { if (head >= pc.heads) {
return; return;
} }
const uint lane = gl_SubgroupInvocationID; const uint lane = gl_SubgroupInvocationID;
const uint batches = pc.pairs / (64 * pc.heads * pc.tokens); const uint batches = pc.pairs / (64 * pc.heads * pc.tokens);
for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) { for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) {
for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) { for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) {
const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32], vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32],
data_x[base + lane + 64], data_x[base + lane + 96]); data_x[base + lane + 64], data_x[base + lane + 96]);
precise vec4 squares = x * x; precise vec4 squares = x * x;
precise float sum = (squares.x + squares.z) + (squares.y + squares.w); precise float sum = (squares.x + squares.z) + (squares.y + squares.w);
[[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) { [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) {
sum += subgroupShuffle(sum, lane ^ stride); sum += subgroupShuffle(sum, lane ^ stride);
} }
sum = subgroupShuffle(sum, 0); sum = subgroupShuffle(sum, 0);
const float scale = inversesqrt(sum / 128.0 + pc.eps); const float scale = inversesqrt(sum / 128.0 + pc.eps);
precise vec4 weighted = scale * x; precise vec4 weighted = scale * x;
[[unroll]] for (uint i = 0; i < 4; ++i) { [[unroll]] for (uint i = 0; i < 4; ++i) {
weighted[i] *= data_w[pc.w_offset + lane + i * 32]; weighted[i] *= data_w[pc.w_offset + lane + i * 32];
} }
[[unroll]] for (uint i = 0; i < 4; ++i) { [[unroll]] for (uint i = 0; i < 4; ++i) {
const uint col = lane + i * 32; const uint col = lane + i * 32;
precise float value = weighted[i]; precise float value = weighted[i];
if (pc.interleaved != 0) { if (pc.interleaved != 0) {
const float a = subgroupShuffle(value, lane & ~1u); const float a = subgroupShuffle(value, lane & ~1u);
const float b = subgroupShuffle(value, lane | 1u); const float b = subgroupShuffle(value, lane | 1u);
if (col / 2 < pc.rot_half) { if (col / 2 < pc.rot_half) {
const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1]; const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1];
value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]]; value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]];
} }
} else { } else {
const uint pair = col % pc.rot_half; const uint pair = col % pc.rot_half;
const float a = shuffle_channel(weighted, pair); const float a = shuffle_channel(weighted, pair);
const float b = shuffle_channel(weighted, pair + pc.rot_half); const float b = shuffle_channel(weighted, pair + pc.rot_half);
if (col < 2 * pc.rot_half) { if (col < 2 * pc.rot_half) {
const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3]; const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3];
const float c = data_pe[matrix]; const float c = data_pe[matrix];
const float s = data_pe[matrix + pc.pe_stride[1]]; const float s = data_pe[matrix + pc.pe_stride[1]];
value = col < pc.rot_half ? a * c - b * s : a * s + b * c; value = col < pc.rot_half ? a * c - b * s : a * s + b * c;
} }
} }
#ifdef CAST_F16 #ifdef CAST_F16
src/core/vulkan/shaders/rope_apply.comp
+30 -23
View File
@@ -2,9 +2,15 @@
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer PE { float data_pe[]; }; float data_x[];
layout(binding = 2) writeonly buffer D { float data_d[]; }; };
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
layout(binding = 2) writeonly buffer D {
float data_d[];
};
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint half_dim; uint half_dim;
@@ -18,41 +24,42 @@ layout(push_constant) uniform parameter {
uint pe_offset; uint pe_offset;
uint dst_offset; uint dst_offset;
uint interleaved; uint interleaved;
} pc; }
pc;
void main() { void main() {
const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x*gl_WorkGroupSize.x*gl_GlobalInvocationID.y; const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y;
if (index >= pc.pairs) { if (index >= pc.pairs) {
return; return;
} }
const uint j = index % pc.half_dim; const uint j = index % pc.half_dim;
const uint row = index / pc.half_dim; const uint row = index / pc.half_dim;
const uint token = row % pc.tokens; const uint token = row % pc.tokens;
const uint head = (row / pc.tokens) % pc.heads; const uint head = (row / pc.tokens) % pc.heads;
const uint batch = row / (pc.tokens*pc.heads); const uint batch = row / (pc.tokens * pc.heads);
const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j; const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j;
const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j + 1 : j + pc.rot_half; const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half;
const uint x_row = pc.x_offset + head*pc.x_stride[1] + token*pc.x_stride[2] + batch*pc.x_stride[3]; const uint x_row = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
const uint out_row = pc.dst_offset + row*pc.half_dim*2; const uint out_row = pc.dst_offset + row * pc.half_dim * 2;
const float a = data_x[x_row + i0*pc.x_stride[0]]; const float a = data_x[x_row + i0 * pc.x_stride[0]];
const float b = data_x[x_row + i1*pc.x_stride[0]]; const float b = data_x[x_row + i1 * pc.x_stride[0]];
if (j >= pc.rot_half) { if (j >= pc.rot_half) {
data_d[out_row + i0] = a; data_d[out_row + i0] = a;
data_d[out_row + i1] = b; data_d[out_row + i1] = b;
return; return;
} }
const uint matrix = pc.pe_offset + j*pc.pe_stride[2] + token*pc.pe_stride[3]; const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3];
const float c = data_pe[matrix]; const float c = data_pe[matrix];
const float s = data_pe[matrix + pc.pe_stride[1]]; const float s = data_pe[matrix + pc.pe_stride[1]];
// Preserve the separate multiply/add rounding of the unfused graph. // Preserve the separate multiply/add rounding of the unfused graph.
precise float y0; precise float y0;
precise float y1; precise float y1;
if (pc.interleaved != 0) { if (pc.interleaved != 0) {
y0 = a*c + b*data_pe[matrix + pc.pe_stride[0]]; y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]];
y1 = a*s + b*data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]]; y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
} else { } else {
y0 = a*c - b*s; y0 = a * c - b * s;
y1 = a*s + b*c; y1 = a * s + b * c;
} }
data_d[out_row + i0] = y0; data_d[out_row + i0] = y0;
data_d[out_row + i1] = y1; data_d[out_row + i1] = y1;