mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +00:00
refactor: move model-specific operators and fusions into backend extensions (#2253)
This commit is contained in:
leejet authored and GitHub committed
1 parent
c6dfa9def1
commit
45a63ac460
68 files changed
+12173
-151
No files matched your search
@@ -0,0 +1,84 @@
|
||||
cmake_policy(SET CMP0079 NEW)
|
||||
|
||||
function(sd_cuda_sources VENDOR LANGUAGE)
|
||||
set(SOURCES)
|
||||
foreach(OP rope int8 convrot affine registration)
|
||||
set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
|
||||
string(TOUPPER "${VENDOR}" UPPER_VENDOR)
|
||||
file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
|
||||
set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE})
|
||||
list(APPEND SOURCES "${SOURCE}")
|
||||
endforeach()
|
||||
set(TARGET "sd-${VENDOR}")
|
||||
add_library(${TARGET} OBJECT ${SOURCES})
|
||||
target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
|
||||
target_link_libraries(${TARGET} PRIVATE ggml)
|
||||
set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_compile_features(${TARGET} PRIVATE cxx_std_17)
|
||||
target_link_libraries(${SD_LIB} PRIVATE ${TARGET})
|
||||
set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE)
|
||||
set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE)
|
||||
endfunction()
|
||||
|
||||
if(SD_CUDA)
|
||||
enable_language(CUDA)
|
||||
sd_cuda_sources(cuda CUDA)
|
||||
target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu sol_attention.cu)
|
||||
set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
|
||||
set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
|
||||
if(SD_CUDA_CUTLASS_PATH)
|
||||
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
|
||||
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
|
||||
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda")
|
||||
endif()
|
||||
set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas)
|
||||
target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
|
||||
endif()
|
||||
|
||||
if(SD_HIPBLAS)
|
||||
list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr)
|
||||
find_package(hip REQUIRED)
|
||||
find_package(hipblas REQUIRED)
|
||||
if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc")
|
||||
sd_cuda_sources(hip CXX)
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device)
|
||||
else()
|
||||
sd_cuda_sources(hip HIP)
|
||||
set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
|
||||
endif()
|
||||
set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
|
||||
target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
|
||||
endif()
|
||||
|
||||
if(SD_MUSA)
|
||||
if(NOT MUSA_PATH)
|
||||
if(DEFINED ENV{MUSA_PATH})
|
||||
set(MUSA_PATH "$ENV{MUSA_PATH}")
|
||||
elseif(EXISTS /opt/musa)
|
||||
set(MUSA_PATH /opt/musa)
|
||||
else()
|
||||
set(MUSA_PATH /usr/local/musa)
|
||||
endif()
|
||||
endif()
|
||||
list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
|
||||
find_package(MUSAToolkit REQUIRED)
|
||||
set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++")
|
||||
sd_cuda_sources(musa CXX)
|
||||
set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero)
|
||||
if(NOT DEFINED MUSA_ARCHITECTURES)
|
||||
set(MUSA_ARCHITECTURES "22;31")
|
||||
endif()
|
||||
foreach(ARCH ${MUSA_ARCHITECTURES})
|
||||
list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}")
|
||||
endforeach()
|
||||
set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}")
|
||||
if(GGML_STATIC)
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static)
|
||||
target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static)
|
||||
else()
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas)
|
||||
target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas)
|
||||
endif()
|
||||
endif()
|
||||
@@ -0,0 +1,183 @@
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include "ggml-ext.h"
|
||||
#include "runtime.h"
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
namespace sd_affine_cuda {
|
||||
static bool disabled() {
|
||||
static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
|
||||
return value;
|
||||
}
|
||||
static bool aligned(const ggml_tensor* tensor, size_t alignment) {
|
||||
size_t offset = 0;
|
||||
while (!tensor->data && tensor->view_src) {
|
||||
offset += tensor->view_offs;
|
||||
tensor = tensor->view_src;
|
||||
}
|
||||
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
|
||||
}
|
||||
static float param(const ggml_tensor* tensor, int index) {
|
||||
float value;
|
||||
std::memcpy(&value, tensor->op_params + index, sizeof(value));
|
||||
return value;
|
||||
}
|
||||
static bool contiguous_f32(const ggml_tensor* tensor) {
|
||||
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
|
||||
}
|
||||
static bool bounded(const ggml_tensor* tensor) {
|
||||
return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
|
||||
}
|
||||
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
|
||||
auto* g = const_cast<ggml_cgraph*>(graph);
|
||||
if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
|
||||
return 0;
|
||||
const auto* node = ggml_graph_node(g, first);
|
||||
const auto* next = ggml_graph_node(g, first + 1);
|
||||
if (node->op == GGML_OP_MUL) {
|
||||
if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node))
|
||||
return 0;
|
||||
const auto* x = node->src[0];
|
||||
const auto* gate = node->src[1];
|
||||
const auto* residual = next->src[0] == node ? next->src[1] : next->src[0];
|
||||
return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) &&
|
||||
node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) &&
|
||||
ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] &&
|
||||
x->ne[0] % 4 == 0 && bounded(next)
|
||||
? 2
|
||||
: 0;
|
||||
}
|
||||
if (node->op == GGML_OP_CONCAT) {
|
||||
const bool permuted = next->op == GGML_OP_PERMUTE;
|
||||
const int count = permuted ? 3 : 2;
|
||||
if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 ||
|
||||
node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) ||
|
||||
!contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1]))
|
||||
return 0;
|
||||
if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] ||
|
||||
next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] ||
|
||||
next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3]))
|
||||
return 0;
|
||||
const auto* dst = ggml_graph_node(g, first + count - 1);
|
||||
return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 &&
|
||||
ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst)
|
||||
? count
|
||||
: 0;
|
||||
}
|
||||
if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) {
|
||||
const auto* src = node->src[0];
|
||||
const auto* permute = ggml_graph_node(g, first + 2);
|
||||
const auto* dst = ggml_graph_node(g, first + 3);
|
||||
return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 &&
|
||||
aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 &&
|
||||
next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f &&
|
||||
permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 &&
|
||||
permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 &&
|
||||
dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 &&
|
||||
ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst)
|
||||
? 4
|
||||
: 0;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
|
||||
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
if (i >= count4)
|
||||
return;
|
||||
const float4 a = x[i], g = gate[i % n4], r = residual[i];
|
||||
// The graph performs two FP32 operations; contracting them into FMA changes rounding.
|
||||
dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
|
||||
__fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
|
||||
}
|
||||
template <bool Permuted>
|
||||
static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
|
||||
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
if (i >= count4)
|
||||
return;
|
||||
const int64_t col = i % columns4, row = i / columns4;
|
||||
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
|
||||
const bool from_prefix = token < prefix_tokens;
|
||||
const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
|
||||
const int64_t src_token = from_prefix ? token : token - prefix_tokens;
|
||||
const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head
|
||||
: (batch * heads + head) * src_tokens + src_token;
|
||||
const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col];
|
||||
dst[2 * i] = __floats2half2_rn(value.x, value.y);
|
||||
dst[2 * i + 1] = __floats2half2_rn(value.z, value.w);
|
||||
}
|
||||
static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
|
||||
const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
|
||||
if (index >= count)
|
||||
return;
|
||||
const int64_t col = index % dim, row = index / dim;
|
||||
const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
|
||||
const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
|
||||
dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
|
||||
dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
|
||||
}
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
|
||||
*workspace = ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void destroy(void*, void*) {}
|
||||
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<CUDAStream>(context.stream);
|
||||
const auto* node = nodes[0];
|
||||
const auto* a = node->src[0];
|
||||
const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
|
||||
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 bool gate = residual != nullptr;
|
||||
const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
|
||||
(residual && overlaps(dst, residual) && dst->data != residual->data);
|
||||
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
|
||||
const int64_t count4 = ggml_nelements(dst) / 4;
|
||||
const unsigned blocks = unsigned((count4 + 255) / 256);
|
||||
if (gate) {
|
||||
gated_residual<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
|
||||
static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
|
||||
} else if (node->op == GGML_OP_CONCAT) {
|
||||
const bool permuted = count == 3;
|
||||
auto kernel = permuted ? concat_cast<true> : concat_cast<false>;
|
||||
kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
|
||||
static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
|
||||
node->ne[permuted ? 1 : 2], count4);
|
||||
} else {
|
||||
scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
|
||||
a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
|
||||
}
|
||||
auto error = cuda_last_error();
|
||||
if (error == cuda_success && alias)
|
||||
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
|
||||
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
}
|
||||
|
||||
void register_affine(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.priority = 200;
|
||||
desc.match = sd_affine_cuda::match;
|
||||
desc.prepare = sd_affine_cuda::prepare;
|
||||
desc.destroy = sd_affine_cuda::destroy;
|
||||
desc.execute = sd_affine_cuda::execute;
|
||||
desc.name = "sd.gated_residual." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_MUL;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_CONCAT;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_CONT;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include "convrot_kernel.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
#include "runtime.h"
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
namespace sd_convrot {
|
||||
|
||||
static bool fusion_disabled() {
|
||||
static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
|
||||
return disabled;
|
||||
}
|
||||
static float param_f32(const ggml_tensor* tensor, int index) {
|
||||
float value;
|
||||
std::memcpy(&value, tensor->op_params + index, sizeof(value));
|
||||
return value;
|
||||
}
|
||||
static bool f32(const ggml_tensor* tensor) {
|
||||
return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
|
||||
}
|
||||
static bool vector(const ggml_tensor* tensor, int64_t k) {
|
||||
return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
|
||||
}
|
||||
enum Kind { SCALE,
|
||||
NORM,
|
||||
RMS_MOD,
|
||||
SWIGLU,
|
||||
SILU,
|
||||
SIGMOID };
|
||||
struct Operation {
|
||||
Kind kind;
|
||||
const ggml_tensor* input = nullptr;
|
||||
const ggml_tensor* up = nullptr;
|
||||
const ggml_tensor* modulation = nullptr;
|
||||
const ggml_tensor* weight = nullptr;
|
||||
const ggml_tensor* shift = nullptr;
|
||||
int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0;
|
||||
float eps = 0.f, scale = 1.f, bias = 0.f;
|
||||
};
|
||||
static bool describe(ggml_tensor* const* nodes, int count, Operation& op) {
|
||||
int group;
|
||||
const auto* first = nodes[0];
|
||||
const auto* quant = nodes[count - 1];
|
||||
if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
|
||||
return false;
|
||||
const auto* value = quant->src[0];
|
||||
op.k = value->ne[0];
|
||||
op.rows = ggml_nrows(value);
|
||||
op.input = first->src[0];
|
||||
if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX)
|
||||
return false;
|
||||
op.stride = op.input->ne[0];
|
||||
op.up_stride = op.k;
|
||||
if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
|
||||
op.kind = SCALE;
|
||||
op.scale = param_f32(first, 0);
|
||||
op.bias = param_f32(first, 1);
|
||||
} else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
|
||||
op.kind = NORM;
|
||||
op.modulation = nodes[1]->src[1];
|
||||
op.eps = param_f32(first, 0);
|
||||
if (op.k < 1024 || !vector(op.modulation, op.k))
|
||||
return false;
|
||||
} else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
|
||||
nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] &&
|
||||
nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] &&
|
||||
nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) {
|
||||
op.kind = RMS_MOD;
|
||||
op.weight = nodes[1]->src[1];
|
||||
op.modulation = nodes[2]->src[1];
|
||||
op.shift = nodes[4]->src[1];
|
||||
op.eps = param_f32(first, 0);
|
||||
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
|
||||
return false;
|
||||
} else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
|
||||
op.kind = SWIGLU;
|
||||
op.up = first->src[1] ? first->src[1] : op.input;
|
||||
if (!f32(op.up))
|
||||
return false;
|
||||
op.up_stride = op.up->ne[0];
|
||||
if (!first->src[1]) {
|
||||
if (first->op_params[1])
|
||||
op.gate_offset = op.k;
|
||||
else
|
||||
op.up_offset = op.k;
|
||||
}
|
||||
} else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] &&
|
||||
(nodes[1]->src[0] == first || nodes[1]->src[1] == first)) {
|
||||
auto unary = ggml_get_unary_op(first);
|
||||
if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID)
|
||||
return false;
|
||||
op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID;
|
||||
op.up = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0];
|
||||
if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up))
|
||||
return false;
|
||||
} else
|
||||
return false;
|
||||
for (int i = 0; i < count - 1; ++i) {
|
||||
if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value))
|
||||
return false;
|
||||
for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
|
||||
if (input == nodes[i])
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) {
|
||||
const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
|
||||
: 128;
|
||||
return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory;
|
||||
}
|
||||
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
|
||||
if (fusion_disabled())
|
||||
return 0;
|
||||
auto* g = const_cast<ggml_cgraph*>(graph);
|
||||
const int available = ggml_graph_n_nodes(g) - first;
|
||||
const auto* node = ggml_graph_node(g, first);
|
||||
const int count = node->op == GGML_OP_RMS_NORM ? 6 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3
|
||||
: 2;
|
||||
if (available < count)
|
||||
return 0;
|
||||
ggml_tensor* nodes[6];
|
||||
for (int j = 0; j < count; ++j)
|
||||
nodes[j] = ggml_graph_node(g, first + j);
|
||||
Operation op;
|
||||
return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
|
||||
}
|
||||
static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
|
||||
if (fusion_disabled())
|
||||
return;
|
||||
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
|
||||
auto* quant = ggml_graph_node(graph, i);
|
||||
int group;
|
||||
if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
|
||||
continue;
|
||||
ggml_tensor* chain[6] = {};
|
||||
int count = 0;
|
||||
auto* value = quant->src[0];
|
||||
if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) {
|
||||
auto* add = value->src[0];
|
||||
auto* weight = add->src[0];
|
||||
if (weight->op == GGML_OP_MUL) {
|
||||
chain[0] = weight->src[0];
|
||||
chain[1] = weight;
|
||||
chain[2] = add->src[1];
|
||||
chain[3] = add;
|
||||
chain[4] = value;
|
||||
chain[5] = quant;
|
||||
count = 6;
|
||||
}
|
||||
} else if (value->op == GGML_OP_MUL) {
|
||||
for (int src = 0; src < 2; ++src)
|
||||
if (value->src[src]->op == GGML_OP_UNARY) {
|
||||
chain[0] = value->src[src];
|
||||
chain[1] = value;
|
||||
chain[2] = quant;
|
||||
count = 3;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Operation op;
|
||||
if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
|
||||
continue;
|
||||
int begin = -1;
|
||||
for (int j = std::max(0, i - 64); j < i; ++j)
|
||||
if (ggml_graph_node(graph, j) == chain[0]) {
|
||||
begin = j;
|
||||
break;
|
||||
}
|
||||
if (begin < 0)
|
||||
continue;
|
||||
// Move the owning producer with its in-place consumers so alias checks retain their storage relationship.
|
||||
for (int part = 1; part < count; ++part) {
|
||||
int next = -1;
|
||||
for (int j = begin + part; j <= i; ++j)
|
||||
if (ggml_graph_node(graph, j) == chain[part]) {
|
||||
next = j;
|
||||
break;
|
||||
}
|
||||
if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next)))
|
||||
break;
|
||||
begin = next - part;
|
||||
}
|
||||
}
|
||||
}
|
||||
struct Plan {
|
||||
ggml_ext_gpu_device_info device;
|
||||
Operation op;
|
||||
bool configured = false;
|
||||
};
|
||||
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
|
||||
auto* plan = new Plan{};
|
||||
*state = plan;
|
||||
plan->device = cuda_device_info(ggml_backend_get_device(backend));
|
||||
if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
|
||||
return GGML_STATUS_FAILED;
|
||||
*workspace = ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void destroy(void* state, void*) {
|
||||
delete static_cast<Plan*>(state);
|
||||
}
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
if (!b)
|
||||
return false;
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
|
||||
static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
|
||||
constexpr bool warp = Scaled || RMS;
|
||||
constexpr int extra = Normalize ? 384 : RMS ? 256
|
||||
: 128;
|
||||
if (!p.configured) {
|
||||
auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
|
||||
int(p.device.shared_memory - extra));
|
||||
if (status != cuda_success)
|
||||
return status;
|
||||
p.configured = true;
|
||||
}
|
||||
const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
|
||||
: op.k <= 8192 ? 512
|
||||
: 1024)
|
||||
: Swiglu && !Sigmoid && op.k >= 16384 ? 512
|
||||
: 256;
|
||||
const int64_t padded = (op.rows + 3) / 4 * 4;
|
||||
auto data = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
|
||||
quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
|
||||
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
|
||||
op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
|
||||
return cuda_last_error();
|
||||
}
|
||||
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
|
||||
auto& plan = *static_cast<Plan*>(state);
|
||||
const auto& op = plan.op;
|
||||
auto* dst = nodes[count - 1];
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<CUDAStream>(context.stream);
|
||||
bool alias = false;
|
||||
for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
|
||||
alias |= overlaps(dst, input);
|
||||
auto* output = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
|
||||
const int64_t padded = (op.rows + 3) / 4 * 4;
|
||||
if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
|
||||
return GGML_STATUS_FAILED;
|
||||
CUDAError error;
|
||||
switch (op.kind) {
|
||||
case SCALE:
|
||||
error = launch<false, false, true>(plan, op, output, stream);
|
||||
break;
|
||||
case NORM:
|
||||
error = launch<false, true>(plan, op, output, stream);
|
||||
break;
|
||||
case RMS_MOD:
|
||||
error = launch<false, false, false, false, true>(plan, op, output, stream);
|
||||
break;
|
||||
case SWIGLU:
|
||||
case SILU:
|
||||
error = launch<true>(plan, op, output, stream);
|
||||
break;
|
||||
case SIGMOID:
|
||||
error = launch<true, false, false, true>(plan, op, output, stream);
|
||||
break;
|
||||
default:
|
||||
return GGML_STATUS_FAILED;
|
||||
}
|
||||
if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
|
||||
return GGML_STATUS_FAILED;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
}
|
||||
|
||||
void register_convrot(ggml_backend_dev_t device) {
|
||||
for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
const std::string name = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
|
||||
desc.name = name.c_str();
|
||||
desc.op = op;
|
||||
desc.priority = 225;
|
||||
desc.optimize = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
|
||||
desc.match = sd_convrot::match;
|
||||
desc.prepare = sd_convrot::prepare;
|
||||
desc.execute = sd_convrot::execute;
|
||||
desc.destroy = sd_convrot::destroy;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
|
||||
#define __SD_CORE_CUDA_CONVROT_KERNEL_H__
|
||||
#include <cstdint>
|
||||
#include "runtime.h"
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
namespace sd_convrot {
|
||||
static __device__ __forceinline__ float warp_reduce_sum(float x) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
x += cuda_shuffle_xor(x, offset, 32);
|
||||
return x;
|
||||
}
|
||||
static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
x.x += cuda_shuffle_xor(x.x, offset, 32);
|
||||
x.y += cuda_shuffle_xor(x.y, offset, 32);
|
||||
}
|
||||
return x;
|
||||
}
|
||||
template <int Width>
|
||||
static __device__ __forceinline__ float warp_reduce_max(float x) {
|
||||
#pragma unroll
|
||||
for (int offset = Width / 2; offset > 0; offset >>= 1)
|
||||
x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
|
||||
return x;
|
||||
}
|
||||
static __device__ float2 block_sum(float2 x, float2* shared) {
|
||||
x = warp_reduce_sum(x);
|
||||
if (threadIdx.x % 32 == 0)
|
||||
shared[threadIdx.x / 32] = x;
|
||||
__syncthreads();
|
||||
return warp_reduce_sum(shared[threadIdx.x % 32]);
|
||||
}
|
||||
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);
|
||||
}
|
||||
|
||||
template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
|
||||
__global__ void quantize(
|
||||
const T* src,
|
||||
int8_t* dst,
|
||||
float* scales,
|
||||
int64_t k,
|
||||
int64_t rows,
|
||||
int64_t dst_stride,
|
||||
const float* up,
|
||||
int64_t src_stride,
|
||||
int64_t up_stride,
|
||||
const float* activation_scales = nullptr,
|
||||
const float* weight_scales = nullptr,
|
||||
const float* modulation = nullptr,
|
||||
float eps = 0.0f,
|
||||
float input_scale = 1.0f,
|
||||
float input_bias = 0.0f,
|
||||
const float* rms_weight = nullptr,
|
||||
const float* rms_shift = nullptr) {
|
||||
constexpr int group_size = 256;
|
||||
|
||||
const int64_t row = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (row >= rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
extern __shared__ float values[];
|
||||
__shared__ float maxima[32];
|
||||
const T* src_row = src + row * src_stride;
|
||||
int8_t* dst_row = dst + row * dst_stride;
|
||||
|
||||
[[maybe_unused]] float rms_scale = 1.0f;
|
||||
if constexpr (rms_modulation) {
|
||||
__shared__ float sums[32];
|
||||
// Preserve the unfused RMSNorm's 1024-thread reduction order with fewer physical warps.
|
||||
for (int virtual_tid = tid; virtual_tid < 1024; virtual_tid += blockDim.x) {
|
||||
float partial = 0.0f;
|
||||
for (int64_t i = virtual_tid; i < k; i += 1024) {
|
||||
const float value = src_row[i];
|
||||
partial += value * value;
|
||||
}
|
||||
partial = warp_reduce_sum(partial);
|
||||
if (tid % 32 == 0) {
|
||||
sums[virtual_tid / 32] = partial;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const float sum = warp_reduce_sum(sums[tid % 32]);
|
||||
rms_scale = rsqrtf(sum / k + eps);
|
||||
}
|
||||
|
||||
if constexpr (swiglu || warp_rotation) {
|
||||
const int lane = tid % 32;
|
||||
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 int64_t i = group * group_size + j * 32 + lane;
|
||||
if constexpr (!swiglu) {
|
||||
float value = src_row[i];
|
||||
if constexpr (rms_modulation) {
|
||||
value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
|
||||
value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]);
|
||||
}
|
||||
if constexpr (scaled) {
|
||||
value = __fmaf_rn(value, input_scale, input_bias);
|
||||
}
|
||||
v[j] = value * (1.0f / 16.0f);
|
||||
} else {
|
||||
float gate_value, up_value;
|
||||
if constexpr (dequantize) {
|
||||
gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row];
|
||||
up_value = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row];
|
||||
gate_value *= weight_scales[i];
|
||||
up_value *= weight_scales[k + i];
|
||||
} else {
|
||||
gate_value = src_row[i];
|
||||
up_value = up[row * up_stride + i];
|
||||
}
|
||||
const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
|
||||
v[j] = (activated * up_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 = cuda_shuffle(v[j], base, 32);
|
||||
const float b = cuda_shuffle(v[j], base + stride, 32);
|
||||
const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
|
||||
const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
|
||||
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j += 2) {
|
||||
const float a = cuda_shuffle(v[j], lane % 16, 32);
|
||||
const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
|
||||
const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
|
||||
const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
|
||||
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
|
||||
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
|
||||
}
|
||||
#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();
|
||||
} else {
|
||||
constexpr int group_threads = group_size / 4;
|
||||
constexpr int groups_per_wave = 1024 / group_threads;
|
||||
if constexpr (normalize) {
|
||||
__shared__ float2 sums[32];
|
||||
float2 mean_var = make_float2(0.0f, 0.0f);
|
||||
for (int64_t i = tid; i < k; i += 1024) {
|
||||
const float value = src_row[i];
|
||||
mean_var.x += value;
|
||||
mean_var.y += value * value;
|
||||
}
|
||||
mean_var = block_sum(mean_var, sums);
|
||||
const float mean = mean_var.x / k;
|
||||
const float variance = mean_var.y / k - mean * mean;
|
||||
const float inv_std = rsqrtf(variance + eps);
|
||||
for (int64_t i = tid; i < k; i += 1024) {
|
||||
const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std);
|
||||
values[i] = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f);
|
||||
}
|
||||
} else {
|
||||
for (int64_t i = tid; i < k; i += blockDim.x) {
|
||||
const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i];
|
||||
values[i] = value * (1.0f / 16.0f);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int group_lane = tid % group_threads;
|
||||
const int wave_group = tid / group_threads;
|
||||
const int64_t groups = k / group_size;
|
||||
for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) {
|
||||
const int64_t group = group_base + wave_group;
|
||||
#pragma unroll
|
||||
for (int stride = 1; stride < group_size; stride *= 4) {
|
||||
if (group < groups) {
|
||||
float* group_values = values + group * group_size;
|
||||
const int base = (group_lane / stride) * 4 * stride + group_lane % stride;
|
||||
const int i0 = base;
|
||||
const int i1 = i0 + stride;
|
||||
const int i2 = i1 + stride;
|
||||
const int i3 = i2 + stride;
|
||||
const float a = group_values[i0];
|
||||
const float b = group_values[i1];
|
||||
const float c = group_values[i2];
|
||||
const float d = group_values[i3];
|
||||
group_values[i0] = a + b + c - d;
|
||||
group_values[i1] = a + b - c + d;
|
||||
group_values[i2] = a - b + c + d;
|
||||
group_values[i3] = -a + b + c + d;
|
||||
}
|
||||
__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 = warp_reduce_max<32>(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 = warp_reduce_max<32>(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;
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,792 @@
|
||||
/*
|
||||
* 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 "core/cuda/i8_projection.h"
|
||||
#include "core/cuda/rope.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
#include "ggml-ext.h"
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
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 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);
|
||||
int convrot_group;
|
||||
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 (sd_ggml_convrot_describe(node, &convrot_group)) {
|
||||
if (convrot_group == 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 (ggml_ext_op_get(graph_node(graph, first)) == sd_ggml_i8_op()) {
|
||||
if (first + 5 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
|
||||
!sd_cuda_i8_rope_supports(device, graph_node(graph, first), graph_node(graph, first + 1),
|
||||
graph_node(graph, first + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) ||
|
||||
graph_node(graph, first + 5)->src[0] != graph_node(graph, first + 4))
|
||||
return 0;
|
||||
offset = 5;
|
||||
} else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
|
||||
if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
|
||||
!sd_cuda::rope_fusion_supports(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) &&
|
||||
sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
|
||||
if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
|
||||
ggml_ext_graph_delay(graph, j - 2, 5, i))
|
||||
break;
|
||||
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 = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : 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);
|
||||
}
|
||||
|
||||
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 == 5) {
|
||||
if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
} else if (i == 0 && plan.index == 3) {
|
||||
if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
} 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];
|
||||
int convrot_group;
|
||||
if (sd_ggml_convrot_describe(dst, &convrot_group)) {
|
||||
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_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
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_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.name = "sd.i8_rope_cute_attention.cuda.v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_i8_op();
|
||||
desc.priority = 400;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
@@ -0,0 +1,461 @@
|
||||
#include <cuda_runtime.h>
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
|
||||
#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
|
||||
#include <mutex>
|
||||
#include <unordered_map>
|
||||
#include "core/ggml_extensions.h"
|
||||
|
||||
namespace sd_i8_gated {
|
||||
struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
|
||||
using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
|
||||
using Base::Base;
|
||||
|
||||
CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
|
||||
const auto shape = tiled_shape();
|
||||
if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
|
||||
// Column bands retain weight tiles in Ada's L2 across successive row tiles.
|
||||
constexpr int band_width = 16;
|
||||
if (index >= shape.m() * shape.n()) {
|
||||
return {shape.m(), shape.n(), get_batch_idx()};
|
||||
}
|
||||
const int band = index / (shape.m() * band_width);
|
||||
const int width = min(band_width, shape.n() - band * band_width);
|
||||
const int within = index - band * shape.m() * band_width;
|
||||
return {within / width, band * band_width + within % width, get_batch_idx()};
|
||||
}
|
||||
#endif
|
||||
return Base::get_tile_offset(index);
|
||||
}
|
||||
};
|
||||
|
||||
template <int M, int N, int WM>
|
||||
struct TileConfig {
|
||||
using Tile = cutlass::gemm::GemmShape<M, N, 64>;
|
||||
using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
|
||||
};
|
||||
|
||||
enum Tile {
|
||||
TILE_128_256,
|
||||
TILE_128_128,
|
||||
TILE_64_128,
|
||||
TILE_16_128,
|
||||
};
|
||||
|
||||
static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
|
||||
static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
|
||||
return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
|
||||
}
|
||||
|
||||
static Tile select_tile(int64_t rows, int64_t n, int nsm) {
|
||||
if (rows < 32) {
|
||||
return TILE_16_128;
|
||||
}
|
||||
const int64_t columns = (n + 127) / 128;
|
||||
if (((rows + 15) / 16) * columns <= nsm) {
|
||||
return TILE_16_128;
|
||||
}
|
||||
if (((rows + 63) / 64) * columns <= nsm) {
|
||||
return TILE_64_128;
|
||||
}
|
||||
const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
|
||||
const int64_t square_tiles = ((rows + 127) / 128) * columns;
|
||||
// A narrower tile costs more per output element, but can waste less of the last wave.
|
||||
const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
|
||||
const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
|
||||
return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
|
||||
}
|
||||
|
||||
template <class T>
|
||||
struct MulRound {
|
||||
CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
|
||||
T result;
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < T::kElements; ++i) {
|
||||
result[i] = __fmul_rn(a[i], b[i]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
template <class T>
|
||||
struct AddRound {
|
||||
CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
|
||||
T result;
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < T::kElements; ++i) {
|
||||
result[i] = __fadd_rn(a[i], b[i]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
template <class Config, bool scaled, bool stream_k>
|
||||
struct KernelTypes {
|
||||
using Tile = typename Config::Tile;
|
||||
using Warp = typename Config::Warp;
|
||||
using Swizzle = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
|
||||
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
|
||||
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
|
||||
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
|
||||
using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
|
||||
using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
|
||||
using Residual = cutlass::epilogue::threadblock::VisitorAuxLoad<Map, float, cute::Stride<int64_t, cute::_1, int64_t>>;
|
||||
using GateMul = cutlass::epilogue::threadblock::VisitorCompute<MulRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using Add = cutlass::epilogue::threadblock::VisitorCompute<AddRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using Scalar = cutlass::epilogue::threadblock::VisitorScalarBroadcast<float>;
|
||||
using ScaleOut = cutlass::epilogue::threadblock::Sm80EVT<GateMul, ScaleW, Scalar>;
|
||||
using Scaled = std::conditional_t<scaled, ScaleOut, ScaleW>;
|
||||
using Gated = cutlass::epilogue::threadblock::Sm80EVT<GateMul, Scaled, Row>;
|
||||
using WithResidual = cutlass::epilogue::threadblock::Sm80EVT<Add, Gated, Residual>;
|
||||
using Output = WithResidual;
|
||||
using Store = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
|
||||
using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, Output>;
|
||||
using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
|
||||
int8_t,
|
||||
cutlass::layout::RowMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
int8_t,
|
||||
cutlass::layout::ColumnMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
float,
|
||||
cutlass::layout::RowMajor,
|
||||
4,
|
||||
int,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
Tile,
|
||||
Warp,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
Epilogue,
|
||||
Swizzle,
|
||||
4,
|
||||
cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
|
||||
};
|
||||
|
||||
struct DeviceInfo {
|
||||
int ordinal = -1;
|
||||
int major = 0;
|
||||
int sms = 0;
|
||||
int threads_per_sm = 0;
|
||||
size_t shared = 0;
|
||||
};
|
||||
static DeviceInfo device_info(ggml_backend_dev_t device) {
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
auto found = cache.find(device);
|
||||
if (found != cache.end())
|
||||
return found->second;
|
||||
DeviceInfo info;
|
||||
auto reg = ggml_backend_dev_backend_reg(device);
|
||||
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
|
||||
int ordinal = query ? query(device) : -1;
|
||||
cudaDeviceProp props = {};
|
||||
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
|
||||
info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
|
||||
cache.emplace(device, info);
|
||||
return info;
|
||||
}
|
||||
static float param_f32(const ggml_tensor* tensor, int index) {
|
||||
float value;
|
||||
std::memcpy(&value, tensor->op_params + index, sizeof(value));
|
||||
return value;
|
||||
}
|
||||
static bool aligned(const ggml_tensor* tensor) {
|
||||
if (!tensor)
|
||||
return false;
|
||||
size_t offset = 0;
|
||||
while (tensor->view_src) {
|
||||
offset += tensor->view_offs;
|
||||
tensor = tensor->view_src;
|
||||
}
|
||||
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
|
||||
}
|
||||
struct Operation {
|
||||
SDI8Operation mat;
|
||||
const ggml_tensor* gate;
|
||||
const ggml_tensor* residual;
|
||||
const ggml_tensor* prefix;
|
||||
float scale;
|
||||
int64_t rows, n, k;
|
||||
};
|
||||
static Operation describe(ggml_tensor* const* nodes, int count) {
|
||||
Operation op = {};
|
||||
sd_ggml_i8_describe(nodes[0], &op.mat);
|
||||
const bool concat = nodes[count - 1]->op == GGML_OP_CONCAT;
|
||||
const auto* add = nodes[count - (concat ? 2 : 1)];
|
||||
const auto* mul = nodes[count - (concat ? 3 : 2)];
|
||||
op.gate = mul->src[1];
|
||||
op.residual = add->src[0] == mul ? add->src[1] : add->src[0];
|
||||
op.prefix = concat ? nodes[count - 1]->src[0] : nullptr;
|
||||
op.scale = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
|
||||
op.rows = ggml_nrows(nodes[0]);
|
||||
op.n = nodes[0]->ne[0];
|
||||
op.k = op.mat.weight->ne[0];
|
||||
return op;
|
||||
}
|
||||
template <class Types>
|
||||
static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) {
|
||||
const auto* x = metadata ? nullptr : static_cast<const int8_t*>(op.mat.input->data);
|
||||
const auto* w = metadata ? nullptr : static_cast<const int8_t*>(op.mat.weight->data);
|
||||
const auto* xs = metadata ? nullptr : op.mat.packed ? reinterpret_cast<const float*>(x + op.k * ((op.rows + 3) / 4 * 4))
|
||||
: static_cast<const float*>(op.mat.input_scale->data);
|
||||
const auto* ws = metadata ? nullptr : static_cast<const float*>(op.mat.weight_scale->data);
|
||||
const auto* gate = metadata ? nullptr : static_cast<const float*>(op.gate->data);
|
||||
auto* residual = metadata ? nullptr : static_cast<float*>(op.residual->data);
|
||||
typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
|
||||
typename Types::Scaled::Arguments scaled;
|
||||
if constexpr (std::is_same<typename Types::Scaled, typename Types::ScaleW>::value)
|
||||
scaled = scales;
|
||||
else
|
||||
scaled = {scales, {{op.scale}}, {}};
|
||||
typename Types::Epilogue::Arguments epilogue = {{{scaled, {gate, 0.f, {}}, {}}, {residual, 0.f, {op.n, cute::_1{}, op.rows * op.n}}, {}}, {output, {op.n, cute::_1{}, op.rows * op.n}}};
|
||||
return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
|
||||
}
|
||||
struct Plan {
|
||||
DeviceInfo device;
|
||||
size_t kernel_workspace = 0;
|
||||
size_t output_offset = 0;
|
||||
int occupancy = 0;
|
||||
ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
|
||||
};
|
||||
template <class Types>
|
||||
static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) {
|
||||
using Kernel = typename Types::Kernel;
|
||||
auto args = arguments<Types>(op, output, false);
|
||||
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
if (plan.occupancy == 0) {
|
||||
if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
|
||||
cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
int occupancy = 0;
|
||||
if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2<Kernel>,
|
||||
Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
|
||||
occupancy <= 0)
|
||||
return GGML_STATUS_FAILED;
|
||||
plan.occupancy = occupancy;
|
||||
}
|
||||
typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
|
||||
if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
cutlass::arch::synclog_setup();
|
||||
cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
|
||||
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
template <class Config, bool Scaled, bool StreamK>
|
||||
static bool configure(const Operation& op, Plan& plan) {
|
||||
using Types = KernelTypes<Config, Scaled, StreamK>;
|
||||
using Kernel = typename Types::Kernel;
|
||||
if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
|
||||
return false;
|
||||
auto args = arguments<Types>(op, nullptr, true);
|
||||
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
|
||||
return false;
|
||||
size_t workspace = 0;
|
||||
// Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
|
||||
for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
|
||||
typename Kernel::Params params(args, plan.device.sms, occupancy);
|
||||
const auto grid = params.get_grid_dims();
|
||||
if (grid.y > 65535 || grid.z > 65535)
|
||||
return false;
|
||||
workspace = std::max(workspace, params.get_workspace_size());
|
||||
}
|
||||
plan.kernel_workspace = workspace;
|
||||
plan.output_offset = (workspace + 255) / 256 * 256;
|
||||
plan.launch = launch<Types>;
|
||||
return true;
|
||||
}
|
||||
template <class Config>
|
||||
static bool configure(const Operation& op, Plan& plan) {
|
||||
if (op.scale != 1.f)
|
||||
return op.rows >= 8192 ? configure<Config, true, true>(op, plan) : configure<Config, true, false>(op, plan);
|
||||
return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
|
||||
}
|
||||
static bool configure(const Operation& op, Plan& plan) {
|
||||
auto try_tile = [&](Tile tile) {
|
||||
switch (tile) {
|
||||
case TILE_128_256:
|
||||
return configure<TileConfig<128, 256, 64>>(op, plan);
|
||||
case TILE_128_128:
|
||||
return configure<TileConfig<128, 128, 64>>(op, plan);
|
||||
case TILE_64_128:
|
||||
return configure<TileConfig<64, 128, 32>>(op, plan);
|
||||
case TILE_16_128:
|
||||
return configure<TileConfig<16, 128, 16>>(op, plan);
|
||||
}
|
||||
return false;
|
||||
};
|
||||
const auto tile = select_tile(op.rows, op.n, plan.device.sms);
|
||||
if (try_tile(tile))
|
||||
return true;
|
||||
for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
|
||||
if (fallback != tile && try_tile(fallback))
|
||||
return true;
|
||||
return false;
|
||||
}
|
||||
static bool fusion_disabled() {
|
||||
static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
|
||||
return disabled;
|
||||
}
|
||||
static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
|
||||
#ifndef SD_CUDA_FORCE_CUBLAS
|
||||
if (fusion_disabled())
|
||||
return;
|
||||
const int total = ggml_graph_n_nodes(graph);
|
||||
for (int i = 2; i + 1 < total; ++i) {
|
||||
const auto* mul = ggml_graph_node(graph, i);
|
||||
const auto* add = ggml_graph_node(graph, i + 1);
|
||||
if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
|
||||
continue;
|
||||
const auto* producer = mul->src[0];
|
||||
const bool scaled = producer->op == GGML_OP_SCALE;
|
||||
const auto* mat = scaled ? producer->src[0] : producer;
|
||||
SDI8Operation op;
|
||||
if (!sd_ggml_i8_describe(mat, &op) || op.input->type != GGML_TYPE_I8 || !op.weight_scale || op.bias)
|
||||
continue;
|
||||
int scale_index = -1;
|
||||
for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
|
||||
if (scaled && ggml_graph_node(graph, j) == producer)
|
||||
scale_index = j;
|
||||
if (ggml_graph_node(graph, j) != mat)
|
||||
continue;
|
||||
if (scaled) {
|
||||
if (scale_index < 0 || (j + 1 < scale_index && !ggml_ext_graph_delay(graph, j, 1, scale_index)))
|
||||
break;
|
||||
j = scale_index - 1;
|
||||
}
|
||||
ggml_ext_graph_delay(graph, j, scaled ? 2 : 1, i);
|
||||
break;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
|
||||
#ifdef SD_CUDA_FORCE_CUBLAS
|
||||
return 0;
|
||||
#else
|
||||
const int available = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) - first;
|
||||
if (fusion_disabled() || available < 3)
|
||||
return 0;
|
||||
ggml_tensor* nodes[5] = {};
|
||||
for (int i = 0; i < std::min(available, 5); ++i)
|
||||
nodes[i] = ggml_graph_node(const_cast<ggml_cgraph*>(graph), first + i);
|
||||
SDI8Operation mat;
|
||||
if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
|
||||
return 0;
|
||||
const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
|
||||
if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
|
||||
((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
|
||||
return 0;
|
||||
for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
|
||||
if (!aligned(t))
|
||||
return 0;
|
||||
const bool scaled = nodes[1]->op == GGML_OP_SCALE;
|
||||
int count = scaled ? 4 : 3;
|
||||
if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
|
||||
return 0;
|
||||
const auto* mul = nodes[count - 2];
|
||||
const auto* add = nodes[count - 1];
|
||||
if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || mul->src[0] != nodes[count - 3] ||
|
||||
(add->src[0] != mul && add->src[1] != mul))
|
||||
return 0;
|
||||
const auto* gate = mul->src[1];
|
||||
const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0];
|
||||
if (!gate || !residual || gate->ne[0] != n || ggml_nelements(gate) != n || !aligned(gate) || !aligned(residual))
|
||||
return 0;
|
||||
for (const auto* t : {gate, residual, mul, add})
|
||||
if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || (t != gate && !ggml_are_same_shape(nodes[0], t)))
|
||||
return 0;
|
||||
if (available > count) {
|
||||
const auto* concat = nodes[count];
|
||||
if (concat->op == GGML_OP_CONCAT && concat->src[1] == add && concat->op_params[0] == 1 &&
|
||||
concat->type == GGML_TYPE_F32 && concat->ne[2] == 1 && concat->ne[3] == 1 && ggml_is_contiguous(concat) &&
|
||||
concat->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(concat->src[0]))
|
||||
++count;
|
||||
}
|
||||
const auto* prefix = nodes[count - 1]->op == GGML_OP_CONCAT ? nodes[count - 1]->src[0] : nullptr;
|
||||
for (int i = 0; i < count; ++i)
|
||||
if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
|
||||
return 0;
|
||||
Plan plan;
|
||||
plan.device = device_info(device);
|
||||
if (plan.device.major < 8 || plan.device.sms <= 0)
|
||||
return 0;
|
||||
return configure(describe(nodes, count), plan) ? count : 0;
|
||||
#endif
|
||||
}
|
||||
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
|
||||
auto* plan = new Plan;
|
||||
*state = plan;
|
||||
plan->device = device_info(ggml_backend_get_device(backend));
|
||||
if (!configure(describe(nodes, count), *plan))
|
||||
return GGML_STATUS_FAILED;
|
||||
*workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void destroy(void* state, void*) {
|
||||
delete static_cast<Plan*>(state);
|
||||
}
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
if (!b)
|
||||
return false;
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
|
||||
auto& plan = *static_cast<Plan*>(state);
|
||||
const auto op = describe(nodes, count);
|
||||
auto* dst = nodes[count - 1];
|
||||
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
|
||||
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
|
||||
ggml_ext_cuda_context context;
|
||||
if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<cudaStream_t>(context.stream);
|
||||
auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
||||
bool alias = reinterpret_cast<uintptr_t>(dst->data) % 16 != 0;
|
||||
for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
|
||||
if (input == op.residual && !op.prefix && input->data == dst->data)
|
||||
continue;
|
||||
alias |= overlaps(dst, input);
|
||||
}
|
||||
auto* output = alias ? scratch + plan.output_offset : static_cast<char*>(dst->data);
|
||||
const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
|
||||
auto status = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);
|
||||
if (status != GGML_STATUS_SUCCESS)
|
||||
return status;
|
||||
if (prefix_bytes && cudaMemcpyAsync(output, op.prefix->data, prefix_bytes, cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
if (alias && cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
}
|
||||
|
||||
void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.i8_gated_residual.cuda.v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_i8_op();
|
||||
desc.priority = 325;
|
||||
desc.optimize = sd_i8_gated::optimize;
|
||||
desc.match = sd_i8_gated::match;
|
||||
desc.prepare = sd_i8_gated::prepare;
|
||||
desc.execute = sd_i8_gated::execute;
|
||||
desc.destroy = sd_i8_gated::destroy;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
|
||||
#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
|
||||
#include <mutex>
|
||||
#include <unordered_map>
|
||||
#include "i8_gemm.h"
|
||||
|
||||
namespace sd_i8_cutlass {
|
||||
struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
|
||||
using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
|
||||
using Base::Base;
|
||||
|
||||
CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
|
||||
const auto shape = tiled_shape();
|
||||
if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
|
||||
// Column bands retain weight tiles in Ada's L2 across successive row tiles.
|
||||
constexpr int band_width = 16;
|
||||
if (index >= shape.m() * shape.n()) {
|
||||
return {shape.m(), shape.n(), get_batch_idx()};
|
||||
}
|
||||
const int band = index / (shape.m() * band_width);
|
||||
const int width = min(band_width, shape.n() - band * band_width);
|
||||
const int within = index - band * shape.m() * band_width;
|
||||
return {within / width, band * band_width + within % width, get_batch_idx()};
|
||||
}
|
||||
#endif
|
||||
return Base::get_tile_offset(index);
|
||||
}
|
||||
};
|
||||
|
||||
template <int M, int N, int WM>
|
||||
struct TileConfig {
|
||||
using Tile = cutlass::gemm::GemmShape<M, N, 64>;
|
||||
using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
|
||||
};
|
||||
|
||||
enum Tile {
|
||||
TILE_128_256,
|
||||
TILE_128_128,
|
||||
TILE_64_128,
|
||||
TILE_16_128,
|
||||
};
|
||||
|
||||
static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
|
||||
static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
|
||||
return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
|
||||
}
|
||||
|
||||
static Tile select_tile(int64_t rows, int64_t n, int nsm) {
|
||||
if (rows < 32) {
|
||||
return TILE_16_128;
|
||||
}
|
||||
const int64_t columns = (n + 127) / 128;
|
||||
if (((rows + 15) / 16) * columns <= nsm) {
|
||||
return TILE_16_128;
|
||||
}
|
||||
if (((rows + 63) / 64) * columns <= nsm) {
|
||||
return TILE_64_128;
|
||||
}
|
||||
const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
|
||||
const int64_t square_tiles = ((rows + 127) / 128) * columns;
|
||||
// A narrower tile costs more per output element, but can waste less of the last wave.
|
||||
const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
|
||||
const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
|
||||
return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
|
||||
}
|
||||
|
||||
struct DeviceInfo {
|
||||
int ordinal = -1;
|
||||
int major = 0;
|
||||
int sms = 0;
|
||||
int threads_per_sm = 0;
|
||||
size_t shared = 0;
|
||||
};
|
||||
static DeviceInfo device_info(ggml_backend_dev_t device) {
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
auto found = cache.find(device);
|
||||
if (found != cache.end())
|
||||
return found->second;
|
||||
DeviceInfo info;
|
||||
auto reg = ggml_backend_dev_backend_reg(device);
|
||||
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
|
||||
int ordinal = query ? query(device) : -1;
|
||||
cudaDeviceProp props = {};
|
||||
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
|
||||
info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
|
||||
cache.emplace(device, info);
|
||||
return info;
|
||||
}
|
||||
// The iterator pairs gate/up columns while reading the original checkpoint layout.
|
||||
struct WeightLayout : cutlass::layout::PermuteBase {
|
||||
int stride, half_n;
|
||||
CUTLASS_HOST_DEVICE WeightLayout(cutlass::PitchLinearCoord extent, int64_t ld)
|
||||
: stride(static_cast<int>(ld)), half_n(extent.strided() / 2) {}
|
||||
CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld)
|
||||
: stride(static_cast<int>(ld)), half_n(extent.column() / 2) {}
|
||||
CUTLASS_HOST_DEVICE int64_t operator()(cutlass::PitchLinearCoord c) const { return c.contiguous() + int64_t(c.strided() / 2 + (c.strided() % 2) * half_n) * stride; }
|
||||
CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
|
||||
};
|
||||
|
||||
template <class Map>
|
||||
struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
|
||||
using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
|
||||
using Base::Base;
|
||||
template <class CB>
|
||||
struct Callbacks : CB {
|
||||
CUTLASS_DEVICE Callbacks(CB const& base)
|
||||
: CB(base) {}
|
||||
CUTLASS_DEVICE void end_step(int step) {
|
||||
auto src = cute::filter(this->tC_rAux);
|
||||
auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < cute::size(src); ++i) {
|
||||
auto coord = coords(i);
|
||||
if (cute::elem_less(coord, this->problem_shape)) {
|
||||
// Each vector contains two gate/up pairs; only their SwiGLU results are stored.
|
||||
const float4 a = *reinterpret_cast<const float4*>(&src(i));
|
||||
float2 value = make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
|
||||
int64_t n = cute::get<1>(this->problem_shape), m = cute::get<0>(this->problem_shape);
|
||||
int64_t offset = cute::get<0>(coord) * (n / 2) + cute::get<1>(coord) / 2 + cute::get<2>(coord) * m * (n / 2);
|
||||
*reinterpret_cast<float2*>(this->params_ptr->ptr_aux + offset) = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
template <class Shape>
|
||||
CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
|
||||
auto cb = Base::get_callbacks(tile, tid, shape);
|
||||
return Callbacks<decltype(cb)>(cb);
|
||||
}
|
||||
};
|
||||
|
||||
static __global__ void pair_scales(const float* src, float* dst, int n) {
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n)
|
||||
dst[i] = src[i / 2 + (i % 2) * (n / 2)];
|
||||
}
|
||||
|
||||
template <class Config, bool StreamK>
|
||||
struct DenseTypes {
|
||||
using Tile = typename Config::Tile;
|
||||
using Warp = typename Config::Warp;
|
||||
using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
|
||||
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
|
||||
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
|
||||
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
|
||||
using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
|
||||
using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
|
||||
using Store = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
|
||||
using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
|
||||
using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
|
||||
int8_t,
|
||||
cutlass::layout::RowMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
int8_t,
|
||||
cutlass::layout::ColumnMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
float,
|
||||
cutlass::layout::RowMajor,
|
||||
4,
|
||||
int,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
Tile,
|
||||
Warp,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
Epilogue,
|
||||
Swizzle,
|
||||
4,
|
||||
cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
|
||||
};
|
||||
template <class Config, bool StreamK>
|
||||
struct SwigluTypes {
|
||||
using Tile = typename Config::Tile;
|
||||
using Warp = typename Config::Warp;
|
||||
using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
|
||||
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
|
||||
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
|
||||
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
|
||||
using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
|
||||
using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
|
||||
using Store = SwigluStore<Map>;
|
||||
using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
|
||||
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 Kernel = std::conditional_t<StreamK,
|
||||
cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
|
||||
cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
|
||||
};
|
||||
struct Plan {
|
||||
DeviceInfo device;
|
||||
int64_t rows, n, k;
|
||||
bool swiglu;
|
||||
size_t kernel_bytes = 0, scratch_offset = 0;
|
||||
int occupancy = 0;
|
||||
ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t) = nullptr;
|
||||
};
|
||||
template <class Types>
|
||||
static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
|
||||
typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
|
||||
typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
|
||||
return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, w, nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
|
||||
}
|
||||
template <class Types>
|
||||
static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
|
||||
using Kernel = typename Types::Kernel;
|
||||
if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
|
||||
reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out)) %
|
||||
16)
|
||||
return GGML_STATUS_FAILED;
|
||||
auto* paired = static_cast<float*>(scratch);
|
||||
auto args = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
|
||||
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
if (!p.occupancy) {
|
||||
if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
|
||||
cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2<Kernel>, Kernel::kThreadCount,
|
||||
sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
|
||||
p.occupancy <= 0)
|
||||
return GGML_STATUS_FAILED;
|
||||
}
|
||||
typename Kernel::Params params(args, p.device.sms, p.occupancy);
|
||||
if (params.get_workspace_size() > p.kernel_bytes)
|
||||
return GGML_STATUS_FAILED;
|
||||
if (p.swiglu) {
|
||||
pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, paired, int(p.n));
|
||||
if (cudaGetLastError() != cudaSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
}
|
||||
if (params.init_workspace(static_cast<char*>(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess)
|
||||
return GGML_STATUS_FAILED;
|
||||
cutlass::arch::synclog_setup();
|
||||
cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
|
||||
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
template <class Types>
|
||||
static bool configure_kernel(Plan& p) {
|
||||
using Kernel = typename Types::Kernel;
|
||||
if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
|
||||
return false;
|
||||
auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
|
||||
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
|
||||
return false;
|
||||
size_t bytes = 0;
|
||||
for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
|
||||
typename Kernel::Params params(args, p.device.sms, occupancy);
|
||||
const auto grid = params.get_grid_dims();
|
||||
if (grid.y > 65535 || grid.z > 65535)
|
||||
return false;
|
||||
bytes = std::max(bytes, params.get_workspace_size());
|
||||
}
|
||||
p.kernel_bytes = bytes;
|
||||
p.scratch_offset = p.swiglu ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
|
||||
p.launch = launch<Types>;
|
||||
return true;
|
||||
}
|
||||
template <class Config>
|
||||
static bool configure_tile(Plan& p) {
|
||||
if (p.swiglu)
|
||||
return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
|
||||
return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
|
||||
}
|
||||
static bool configure(Plan& p) {
|
||||
auto try_tile = [&](Tile tile) {
|
||||
switch (tile) {
|
||||
case TILE_128_256:
|
||||
return configure_tile<TileConfig<128, 256, 64>>(p);
|
||||
case TILE_128_128:
|
||||
return configure_tile<TileConfig<128, 128, 64>>(p);
|
||||
case TILE_64_128:
|
||||
return configure_tile<TileConfig<64, 128, 32>>(p);
|
||||
case TILE_16_128:
|
||||
return configure_tile<TileConfig<16, 128, 16>>(p);
|
||||
}
|
||||
return false;
|
||||
};
|
||||
const auto tile = select_tile(p.rows, p.n, p.device.sms);
|
||||
if (try_tile(tile))
|
||||
return true;
|
||||
for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
|
||||
if (fallback != tile && try_tile(fallback))
|
||||
return true;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
|
||||
SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace) {
|
||||
#ifdef SD_CUDA_FORCE_CUBLAS
|
||||
return nullptr;
|
||||
#else
|
||||
const auto info = sd_i8_cutlass::device_info(device);
|
||||
if (info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
|
||||
rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
|
||||
return nullptr;
|
||||
auto* p = new SDI8CudaGemmPlan;
|
||||
p->device = info;
|
||||
p->rows = rows;
|
||||
p->n = n;
|
||||
p->k = k;
|
||||
p->swiglu = swiglu;
|
||||
if (!sd_i8_cutlass::configure(*p)) {
|
||||
delete p;
|
||||
return nullptr;
|
||||
}
|
||||
*workspace = p->scratch_offset + p->kernel_bytes;
|
||||
return p;
|
||||
#endif
|
||||
}
|
||||
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
|
||||
return plan->launch(*plan, x, w, xs, ws, out, scratch, stream);
|
||||
}
|
||||
void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
|
||||
delete plan;
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
#ifndef __SD_CORE_CUDA_I8_GEMM_H__
|
||||
#define __SD_CORE_CUDA_I8_GEMM_H__
|
||||
#include <cuda_runtime.h>
|
||||
#include "ggml-ext.h"
|
||||
struct SDI8CudaGemmPlan;
|
||||
SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace);
|
||||
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
|
||||
void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan);
|
||||
#endif
|
||||
@@ -0,0 +1,394 @@
|
||||
#include "core/cuda/i8_projection.h"
|
||||
#include "core/cuda/rope.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
|
||||
#include <mutex>
|
||||
#include <type_traits>
|
||||
#include <unordered_map>
|
||||
|
||||
static __device__ __forceinline__ float warp_sum8(float value) {
|
||||
#pragma unroll
|
||||
for (int offset = 4; offset > 0; offset >>= 1)
|
||||
value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset, 8));
|
||||
return value;
|
||||
}
|
||||
|
||||
template <class Map, class Out, bool rms = true>
|
||||
struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
|
||||
using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
|
||||
struct Arguments : Base::Arguments {
|
||||
const float* weight;
|
||||
const float* pe;
|
||||
float eps;
|
||||
float scale;
|
||||
int64_t tokens;
|
||||
half* half_output;
|
||||
int64_t prefix_tokens;
|
||||
};
|
||||
using Params = Arguments;
|
||||
using SharedStorage = typename Base::SharedStorage;
|
||||
const Params* extra;
|
||||
template <class Shape>
|
||||
static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) { return args; }
|
||||
CUTLASS_HOST_DEVICE I8RopeStore(Params const& p, SharedStorage const& s)
|
||||
: Base(p, s), extra(&p) {}
|
||||
template <class CB>
|
||||
struct Callbacks : CB {
|
||||
const Params* extra;
|
||||
CUTLASS_DEVICE Callbacks(CB const& base, const Params* p)
|
||||
: CB(base), extra(p) {}
|
||||
CUTLASS_DEVICE float sum32(float4 v) {
|
||||
float a = warp_sum8(__fmul_rn(v.x, v.x));
|
||||
float b = warp_sum8(__fmul_rn(v.y, v.y));
|
||||
float c = warp_sum8(__fmul_rn(v.z, v.z));
|
||||
float d = warp_sum8(__fmul_rn(v.w, v.w));
|
||||
return __fadd_rn(__fadd_rn(a, c), __fadd_rn(b, d));
|
||||
}
|
||||
CUTLASS_DEVICE void end_step(int step) {
|
||||
static_assert(Map::Detail::kAccessWidth == 16);
|
||||
auto src = cute::filter(this->tC_rAux);
|
||||
auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
|
||||
static_assert(cute::size(src) == 4);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < cute::size(src); i += 2) {
|
||||
float4 va = *reinterpret_cast<const float4*>(&src(i)), vb = *reinterpret_cast<const float4*>(&src(i + 1));
|
||||
float norm = 1.0f;
|
||||
if constexpr (rms) {
|
||||
// Preserve the standalone RMSNorm reduction tree across four 32-element groups.
|
||||
float sum = __fadd_rn(sum32(va), sum32(vb));
|
||||
sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 8));
|
||||
norm = rsqrtf(sum / 128.f + extra->eps);
|
||||
}
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
auto coord = coords(i + j);
|
||||
int row = cute::get<0>(coord), col = cute::get<1>(coord);
|
||||
int heads = cute::get<1>(this->problem_shape) / 128;
|
||||
if (cute::elem_less(coord, this->problem_shape)) {
|
||||
float4 value = j ? vb : va;
|
||||
int local = col % 128;
|
||||
int64_t token = row % extra->tokens, batch = row / extra->tokens, head = col / 128;
|
||||
float4 r = value;
|
||||
if constexpr (rms) {
|
||||
float4 weight = ((const float4*)extra->weight)[local / 4];
|
||||
value.x = __fmul_rn(__fmul_rn(norm, value.x), weight.x);
|
||||
value.y = __fmul_rn(__fmul_rn(norm, value.y), weight.y);
|
||||
value.z = __fmul_rn(__fmul_rn(norm, value.z), weight.z);
|
||||
value.w = __fmul_rn(__fmul_rn(norm, value.w), weight.w);
|
||||
float4 p0 = ((const float4*)extra->pe)[token * 64 + local / 2], p1 = ((const float4*)extra->pe)[token * 64 + local / 2 + 1];
|
||||
r.x = __fadd_rn(__fmul_rn(value.x, p0.x), __fmul_rn(value.y, p0.y));
|
||||
r.y = __fadd_rn(__fmul_rn(value.x, p0.z), __fmul_rn(value.y, p0.w));
|
||||
r.z = __fadd_rn(__fmul_rn(value.z, p1.x), __fmul_rn(value.w, p1.y));
|
||||
r.w = __fadd_rn(__fmul_rn(value.z, p1.z), __fmul_rn(value.w, p1.w));
|
||||
}
|
||||
int64_t offset = ((batch * heads + head) * (extra->tokens + extra->prefix_tokens) + token + extra->prefix_tokens) * 128 + local;
|
||||
if constexpr (std::is_same<Out, half>::value) {
|
||||
half2 a = __floats2half2_rn(__fmul_rn(r.x, extra->scale), __fmul_rn(r.y, extra->scale));
|
||||
half2 b = __floats2half2_rn(__fmul_rn(r.z, extra->scale), __fmul_rn(r.w, extra->scale));
|
||||
half2* dst = (half2*)(extra->half_output + offset);
|
||||
dst[0] = a;
|
||||
dst[1] = b;
|
||||
} else {
|
||||
*((float4*)(this->params_ptr->ptr_aux + offset)) = r;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
template <class Shape>
|
||||
CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
|
||||
auto cb = Base::get_callbacks(tile, tid, shape);
|
||||
return Callbacks<decltype(cb)>(cb, extra);
|
||||
}
|
||||
};
|
||||
|
||||
template <class Out, bool RMS = true>
|
||||
struct ProjectionKernel {
|
||||
using Tile = cutlass::gemm::GemmShape<128, 256, 64>;
|
||||
using Warp = cutlass::gemm::GemmShape<64, 64, 64>;
|
||||
using Swizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
|
||||
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
|
||||
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
|
||||
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
|
||||
using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
|
||||
using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
|
||||
using Store = I8RopeStore<Map, Out, RMS>;
|
||||
using EVT = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
|
||||
using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
|
||||
int8_t,
|
||||
cutlass::layout::RowMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
int8_t,
|
||||
cutlass::layout::ColumnMajor,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
16,
|
||||
float,
|
||||
cutlass::layout::RowMajor,
|
||||
4,
|
||||
int,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
Tile,
|
||||
Warp,
|
||||
cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EVT,
|
||||
Swizzle,
|
||||
4,
|
||||
cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
|
||||
};
|
||||
|
||||
static float param_f32(const ggml_tensor* tensor, int index) {
|
||||
float result;
|
||||
std::memcpy(&result, tensor->op_params + index, sizeof(result));
|
||||
return result;
|
||||
}
|
||||
|
||||
static bool device_supports(ggml_backend_dev_t device) {
|
||||
#ifdef SD_CUDA_FORCE_CUBLAS
|
||||
(void)device;
|
||||
return false;
|
||||
#else
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<ggml_backend_dev_t, bool> cache;
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
auto found = cache.find(device);
|
||||
if (found != cache.end())
|
||||
return found->second;
|
||||
auto reg = ggml_backend_dev_backend_reg(device);
|
||||
auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
|
||||
int ordinal = query ? query(device) : -1;
|
||||
cudaDeviceProp props = {};
|
||||
size_t shared = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
|
||||
sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
|
||||
bool supported = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
|
||||
cache.emplace(device, supported);
|
||||
return supported;
|
||||
#endif
|
||||
}
|
||||
|
||||
static bool aligned(const ggml_tensor* tensor) {
|
||||
size_t offset = 0;
|
||||
while (tensor->view_src) {
|
||||
offset += tensor->view_offs;
|
||||
tensor = tensor->view_src;
|
||||
}
|
||||
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
|
||||
}
|
||||
|
||||
static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
|
||||
SDI8Operation op;
|
||||
if (!sd_ggml_i8_describe(mat, &op) || !op.packed || !op.weight_scale || op.bias)
|
||||
return false;
|
||||
const auto* w = op.weight;
|
||||
const auto* x = op.input;
|
||||
const auto* ws = op.weight_scale;
|
||||
if (mat->type != GGML_TYPE_F32 ||
|
||||
w->type != GGML_TYPE_I8 || x->type != GGML_TYPE_I8 || ws->type != GGML_TYPE_F32 ||
|
||||
w->ne[2] != 1 || w->ne[3] != 1 || mat->ne[0] != w->ne[1] || ggml_nelements(ws) != mat->ne[0] ||
|
||||
!ggml_is_contiguous(mat) || !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || !ggml_is_contiguous(ws) ||
|
||||
!aligned(w) || !aligned(x) || !aligned(ws) || !aligned(mat) ||
|
||||
view->op != GGML_OP_RESHAPE || view->src[0] != mat || view->view_offs != 0 || !ggml_is_contiguous(view) ||
|
||||
view->ne[0] != 128 || view->ne[0] * view->ne[1] != mat->ne[0] || ggml_nelements(view) != ggml_nelements(mat))
|
||||
return false;
|
||||
const int64_t rows = ggml_nrows(mat), n = mat->ne[0], k = w->ne[0];
|
||||
if (rows < 4096 || rows > INT_MAX || n < 4096 || n > INT_MAX || k < 4096 || k > INT_MAX || n % 256 || k % 64 || n / 256 > 65535 ||
|
||||
x->ne[0] != k || x->ne[1] != (rows + 3) / 4 * 4 + (rows * 4 + k - 1) / k || x->ne[2] != 1 || x->ne[3] != 1)
|
||||
return false;
|
||||
return device_supports(device);
|
||||
}
|
||||
|
||||
bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
|
||||
if (!projection_supports(device, mat, view) || norm->src[0] != view ||
|
||||
!sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
|
||||
return false;
|
||||
int32_t interleaved;
|
||||
std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
|
||||
return interleaved != 0;
|
||||
}
|
||||
|
||||
template <class Out, bool RMS = true>
|
||||
static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope, void* output, float scale, int64_t prefix_tokens) {
|
||||
using P = ProjectionKernel<Out, RMS>;
|
||||
using Gemm = typename P::Gemm;
|
||||
const int rows = int(ggml_nrows(mat)), n = int(mat->ne[0]), k = int(mat->src[0]->ne[0]);
|
||||
const auto* x = static_cast<const int8_t*>(mat->src[1]->data);
|
||||
const auto* w = static_cast<const int8_t*>(mat->src[0]->data);
|
||||
const auto* xs = reinterpret_cast<const float*>(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4));
|
||||
const auto* ws = static_cast<const float*>(mat->src[2]->data);
|
||||
const auto* nw = RMS ? static_cast<const float*>(mul->src[1]->data) : nullptr;
|
||||
const auto* pe = RMS ? static_cast<const float*>(rope->src[1]->data) : nullptr;
|
||||
// The visitor forms FP32 addresses even when its only writes target the FP16 output.
|
||||
float* address_space = std::is_same<Out, half>::value ? static_cast<float*>(mat->data) : static_cast<float*>(output);
|
||||
if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
|
||||
reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(nw) | reinterpret_cast<uintptr_t>(pe) |
|
||||
reinterpret_cast<uintptr_t>(address_space)) %
|
||||
16 ||
|
||||
reinterpret_cast<uintptr_t>(output) % 8)
|
||||
return cudaErrorInvalidValue;
|
||||
typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
|
||||
{{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
|
||||
typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
|
||||
int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
|
||||
Gemm gemm;
|
||||
if (gemm.can_implement(args) != cutlass::Status::kSuccess)
|
||||
return cudaErrorInvalidValue;
|
||||
if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess || gemm(stream) != cutlass::Status::kSuccess)
|
||||
return cudaErrorUnknown;
|
||||
return cudaGetLastError();
|
||||
}
|
||||
|
||||
cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale) {
|
||||
return projection_launch<half>(stream, nodes[0], nodes[1], nodes[2], nodes[3], nodes[4], output, scale, 0);
|
||||
}
|
||||
|
||||
static const ggml_tensor* node_at(const ggml_cgraph* graph, int index) {
|
||||
return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
|
||||
}
|
||||
|
||||
static int concat_match(const ggml_cgraph* graph, int first, bool key) {
|
||||
const int offset = key ? 5 : 2, count = key ? 7 : 5;
|
||||
if (first + count > ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
|
||||
return 0;
|
||||
const auto* concat = node_at(graph, first + offset);
|
||||
const auto* copy = node_at(graph, first + count - 1);
|
||||
if (concat->op != GGML_OP_CONCAT || concat->op_params[0] != (key ? 1 : 2) || concat->src[1] != node_at(graph, first + offset - 1) ||
|
||||
copy->op != GGML_OP_CPY || copy->src[0] != node_at(graph, first + count - 2) || copy->type != GGML_TYPE_F16 ||
|
||||
!ggml_is_contiguous(copy) || !ggml_are_same_shape(copy, copy->src[0]))
|
||||
return 0;
|
||||
for (int i = 0; i < 2; ++i)
|
||||
if (!concat->src[i] || concat->src[i]->type != GGML_TYPE_F32 || !ggml_is_contiguous(concat->src[i]))
|
||||
return 0;
|
||||
const auto* view = node_at(graph, first + 1);
|
||||
const auto* prefix = concat->src[0];
|
||||
if (prefix->ne[key ? 1 : 2] > INT_MAX - view->ne[2] || (ggml_nelements(prefix) + 255) / 256 > INT_MAX)
|
||||
return 0;
|
||||
if (!key) {
|
||||
const auto* permute = node_at(graph, first + 3);
|
||||
if (permute->op != GGML_OP_PERMUTE || permute->src[0] != concat || permute->op_params[0] != 0 || permute->op_params[1] != 2 ||
|
||||
permute->op_params[2] != 1 || permute->op_params[3] != 3)
|
||||
return 0;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
|
||||
static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION");return value&&std::atoi(value)!=0; }();
|
||||
if (disabled || first + 4 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
|
||||
return 0;
|
||||
const auto* mat = node_at(graph, first);
|
||||
const auto* view = node_at(graph, first + 1);
|
||||
const auto* norm = node_at(graph, first + 2);
|
||||
const bool key = norm->op == GGML_OP_RMS_NORM;
|
||||
if (key) {
|
||||
if (!sd_cuda_i8_rope_supports(device, mat, view, norm, node_at(graph, first + 3), node_at(graph, first + 4)))
|
||||
return 0;
|
||||
if (int count = concat_match(graph, first, true))
|
||||
return count;
|
||||
if (first + 6 < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph))) {
|
||||
const auto* rope = node_at(graph, first + 4);
|
||||
const auto* scale = node_at(graph, first + 5);
|
||||
const auto* copy = node_at(graph, first + 6);
|
||||
if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
|
||||
copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
|
||||
return 7;
|
||||
}
|
||||
return 5;
|
||||
}
|
||||
return projection_supports(device, mat, view) ? concat_match(graph, first, false) : 0;
|
||||
}
|
||||
|
||||
template <bool Aligned>
|
||||
static __global__ void prefix_copy(const float* input, half2* output, int64_t count4, int64_t prefix, int64_t tokens, int64_t heads, bool permuted) {
|
||||
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
if (i >= count4)
|
||||
return;
|
||||
const int64_t col = i % 32, row = i / 32;
|
||||
const int64_t token = row % prefix, head = (row / prefix) % heads, batch = row / (prefix * heads);
|
||||
const int64_t source_row = permuted ? (batch * prefix + token) * heads + head : (batch * heads + head) * prefix + token;
|
||||
const float* source = input + source_row * 128 + col * 4;
|
||||
float4 value;
|
||||
if constexpr (Aligned)
|
||||
value = *reinterpret_cast<const float4*>(source);
|
||||
else
|
||||
value = make_float4(source[0], source[1], source[2], source[3]);
|
||||
const int64_t target = ((batch * heads + head) * (prefix + tokens) + token) * 32 + col;
|
||||
output[2 * target] = __floats2half2_rn(value.x, value.y);
|
||||
output[2 * target + 1] = __floats2half2_rn(value.z, value.w);
|
||||
}
|
||||
|
||||
static ggml_status projection_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
|
||||
*workspace = ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void projection_destroy(void*, void*) {}
|
||||
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
|
||||
static ggml_status projection_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
|
||||
const auto* mat = nodes[0];
|
||||
const auto* view = nodes[1];
|
||||
const bool key = nodes[2]->op == GGML_OP_RMS_NORM;
|
||||
const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
|
||||
const bool joined = concat->op == GGML_OP_CONCAT;
|
||||
auto* dst = nodes[count - 1];
|
||||
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
|
||||
auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
|
||||
ggml_ext_cuda_context context;
|
||||
if (!query || !query(execution->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<cudaStream_t>(context.stream);
|
||||
bool alias = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
|
||||
for (const auto* input : {mat->src[0], mat->src[1], mat->src[2], key ? nodes[3]->src[1] : nullptr, key ? nodes[4]->src[1] : nullptr, joined ? concat->src[0] : nullptr})
|
||||
if (input)
|
||||
alias |= overlaps(dst, input);
|
||||
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
|
||||
int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
|
||||
float scale = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
|
||||
cudaError_t error;
|
||||
if (dst->type == GGML_TYPE_F32)
|
||||
error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);
|
||||
else if (key)
|
||||
error = projection_launch<half>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix);
|
||||
else
|
||||
error = projection_launch<half, false>(stream, mat, view, nullptr, nullptr, nullptr, output, scale, prefix);
|
||||
if (error == cudaSuccess && joined && prefix > 0) {
|
||||
const auto* input = concat->src[0];
|
||||
const int64_t count4 = ggml_nelements(input) / 4;
|
||||
auto kernel = reinterpret_cast<uintptr_t>(input->data) % 16 == 0 ? prefix_copy<true> : prefix_copy<false>;
|
||||
kernel<<<static_cast<unsigned int>((count4 + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(input->data), static_cast<half2*>(output), count4, prefix, view->ne[2], view->ne[1], !key);
|
||||
error = cudaGetLastError();
|
||||
}
|
||||
if (error == cudaSuccess && alias)
|
||||
error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
|
||||
return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
|
||||
void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.i8_projection.cuda.v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_i8_op();
|
||||
desc.priority = 350;
|
||||
desc.match = projection_match;
|
||||
desc.prepare = projection_prepare;
|
||||
desc.execute = projection_execute;
|
||||
desc.destroy = projection_destroy;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
#ifndef __SD_CORE_CUDA_I8_PROJECTION_H__
|
||||
#define __SD_CORE_CUDA_I8_PROJECTION_H__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "ggml-ext.h"
|
||||
|
||||
bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
|
||||
cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale);
|
||||
void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,739 @@
|
||||
#include <cstdlib>
|
||||
#include "convrot_kernel.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
#ifdef SD_CUDA_I8_CUTLASS
|
||||
#include "core/cuda/i8_gemm.h"
|
||||
#else
|
||||
struct SDI8CudaGemmPlan;
|
||||
#endif
|
||||
|
||||
#include "runtime.h"
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
#include <hipblas/hipblas.h>
|
||||
#elif !defined(SD_CUDA_VENDOR_MUSA)
|
||||
#include <cublas_v2.h>
|
||||
#endif
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <mutex>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
namespace sd_i8_cuda {
|
||||
template <int Width = 32>
|
||||
static __device__ __forceinline__ float warp_reduce_max(float value) {
|
||||
#pragma unroll
|
||||
for (int offset = Width / 2; offset > 0; offset >>= 1)
|
||||
value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
|
||||
return value;
|
||||
}
|
||||
static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
|
||||
const int64_t row = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (row >= rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ float maxima[256];
|
||||
const float* src_row = src + row * k;
|
||||
int8_t* dst_row = dst + row * k;
|
||||
float local_max = 0.0f;
|
||||
for (int64_t i = tid; i < k; i += blockDim.x) {
|
||||
local_max = fmaxf(local_max, fabsf(src_row[i]));
|
||||
}
|
||||
maxima[tid] = local_max;
|
||||
__syncthreads();
|
||||
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const float scale = maxima[0] / 127.0f;
|
||||
if (tid == 0) {
|
||||
scales[row] = scale;
|
||||
}
|
||||
if (scale == 0.0f) {
|
||||
for (int64_t i = tid; i < k; i += blockDim.x) {
|
||||
dst_row[i] = 0;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const float inv_scale = 1.0f / scale;
|
||||
for (int64_t i = tid; i < k; i += blockDim.x) {
|
||||
int value = __float2int_rn(src_row[i] * inv_scale);
|
||||
value = max(-127, min(127, value));
|
||||
dst_row[i] = (int8_t)value;
|
||||
}
|
||||
}
|
||||
static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
|
||||
a = index == 3 ? -a : a;
|
||||
b = index == 2 ? -b : b;
|
||||
c = index == 1 ? -c : c;
|
||||
d = index == 0 ? -d : d;
|
||||
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
|
||||
}
|
||||
static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
|
||||
constexpr int group_size = 256;
|
||||
const int64_t row = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (row >= rows)
|
||||
return;
|
||||
extern __shared__ float values[];
|
||||
__shared__ float maxima[32];
|
||||
const float* src_row = src + row * k;
|
||||
int8_t* dst_row = dst + row * dst_stride;
|
||||
{
|
||||
const int lane = tid % 32;
|
||||
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 int64_t i = group * group_size + j * 32 + lane;
|
||||
const float value = src_row[i];
|
||||
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 = cuda_shuffle(v[j], base, 32);
|
||||
const float b = cuda_shuffle(v[j], base + stride, 32);
|
||||
const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
|
||||
const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
|
||||
v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j += 2) {
|
||||
const float a = cuda_shuffle(v[j], lane % 16, 32);
|
||||
const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
|
||||
const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
|
||||
const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
|
||||
v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
|
||||
v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
|
||||
}
|
||||
#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 = warp_reduce_max<32>(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 = warp_reduce_max<32>(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;
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void convrot_group_amax_i8_cuda(
|
||||
const float* src,
|
||||
float* partial_max,
|
||||
int64_t k,
|
||||
int64_t groups,
|
||||
int64_t total_groups) {
|
||||
const int64_t index = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (index >= total_groups) {
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ float values[256];
|
||||
__shared__ float maxima[2];
|
||||
const int64_t row = index / groups;
|
||||
const int64_t group = index - row * groups;
|
||||
const float* input = src + row * k + group * 256;
|
||||
values[tid] = input[tid] * (1.0f / 16.0f);
|
||||
values[tid + 64] = input[tid + 64] * (1.0f / 16.0f);
|
||||
values[tid + 128] = input[tid + 128] * (1.0f / 16.0f);
|
||||
values[tid + 192] = input[tid + 192] * (1.0f / 16.0f);
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int stride = 1; stride < 256; stride *= 4) {
|
||||
const int base = (tid / stride) * 4 * stride + tid % stride;
|
||||
const int i0 = base;
|
||||
const int i1 = i0 + stride;
|
||||
const int i2 = i1 + stride;
|
||||
const int i3 = i2 + stride;
|
||||
const float a = values[i0];
|
||||
const float b = values[i1];
|
||||
const float c = values[i2];
|
||||
const float d = values[i3];
|
||||
values[i0] = a + b + c - d;
|
||||
values[i1] = a + b - c + d;
|
||||
values[i2] = a - b + c + d;
|
||||
values[i3] = -a + b + c + d;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64]));
|
||||
value = fmaxf(value, fabsf(values[tid + 128]));
|
||||
value = fmaxf(value, fabsf(values[tid + 192]));
|
||||
value = warp_reduce_max<32>(value);
|
||||
if ((tid & 31) == 0) {
|
||||
maxima[tid >> 5] = value;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
partial_max[index] = fmaxf(maxima[0], maxima[1]);
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void reduce_convrot_row_amax_i8_cuda(
|
||||
const float* partial_max,
|
||||
float* scales,
|
||||
int64_t groups,
|
||||
int64_t rows) {
|
||||
const int64_t row = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (row >= rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ float maxima[256];
|
||||
float value = 0.0f;
|
||||
for (int64_t group = tid; group < groups; group += blockDim.x) {
|
||||
value = fmaxf(value, partial_max[row * groups + group]);
|
||||
}
|
||||
maxima[tid] = value;
|
||||
__syncthreads();
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (tid == 0) {
|
||||
scales[row] = maxima[0] / 127.0f;
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void convrot_group_quantize_i8_cuda(
|
||||
const float* src,
|
||||
int8_t* dst,
|
||||
const float* scales,
|
||||
int64_t k,
|
||||
int64_t groups,
|
||||
int64_t total_groups,
|
||||
int64_t dst_stride) {
|
||||
const int64_t index = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (index >= total_groups) {
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ float values[256];
|
||||
const int64_t row = index / groups;
|
||||
const int64_t group = index - row * groups;
|
||||
const float* input = src + row * k + group * 256;
|
||||
int8_t* output = dst + row * dst_stride + group * 256;
|
||||
values[tid] = input[tid] * (1.0f / 16.0f);
|
||||
values[tid + 64] = input[tid + 64] * (1.0f / 16.0f);
|
||||
values[tid + 128] = input[tid + 128] * (1.0f / 16.0f);
|
||||
values[tid + 192] = input[tid + 192] * (1.0f / 16.0f);
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int stride = 1; stride < 256; stride *= 4) {
|
||||
const int base = (tid / stride) * 4 * stride + tid % stride;
|
||||
const int i0 = base;
|
||||
const int i1 = i0 + stride;
|
||||
const int i2 = i1 + stride;
|
||||
const int i3 = i2 + stride;
|
||||
const float a = values[i0];
|
||||
const float b = values[i1];
|
||||
const float c = values[i2];
|
||||
const float d = values[i3];
|
||||
values[i0] = a + b + c - d;
|
||||
values[i1] = a + b - c + d;
|
||||
values[i2] = a - b + c + d;
|
||||
values[i3] = -a + b + c + d;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const float row_scale = scales[row];
|
||||
const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
|
||||
const int indices[4] = {tid, tid + 64, tid + 128, tid + 192};
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
int value = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale);
|
||||
value = max(-127, min(127, value));
|
||||
output[indices[i]] = (int8_t)value;
|
||||
}
|
||||
if (group == 0 && tid == 0 && dst_stride != k) {
|
||||
*((float*)(dst + row * dst_stride + k)) = row_scale;
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void dequantize_i32_rows_cuda(
|
||||
const int32_t* src,
|
||||
const float* scales,
|
||||
const float* weight_scales,
|
||||
const float* bias,
|
||||
float* dst,
|
||||
int64_t n,
|
||||
int64_t rows) {
|
||||
const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535;
|
||||
if (output < n && row < rows) {
|
||||
const int64_t index = row * n + output;
|
||||
const float activation_scale = scales[row];
|
||||
float value = (float)src[index] * activation_scale;
|
||||
if (weight_scales != nullptr) {
|
||||
value *= weight_scales[output];
|
||||
}
|
||||
if (bias != nullptr) {
|
||||
value += bias[output];
|
||||
}
|
||||
dst[index] = value;
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void dequantize_i32_rows_vec4_cuda(
|
||||
const int4* src,
|
||||
const float* scales,
|
||||
const float4* weight_scales,
|
||||
const float4* bias,
|
||||
float4* dst,
|
||||
int64_t n4,
|
||||
int64_t rows) {
|
||||
const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535;
|
||||
if (output >= n4 || row >= rows) {
|
||||
return;
|
||||
}
|
||||
const int64_t index = row * n4 + output;
|
||||
const int4 v = src[index];
|
||||
const float scale = scales[row];
|
||||
float4 value = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale);
|
||||
if (weight_scales != nullptr) {
|
||||
const float4 w = weight_scales[output];
|
||||
value.x *= w.x;
|
||||
value.y *= w.y;
|
||||
value.z *= w.z;
|
||||
value.w *= w.w;
|
||||
}
|
||||
if (bias != nullptr) {
|
||||
const float4 b = bias[output];
|
||||
value.x += b.x;
|
||||
value.y += b.y;
|
||||
value.z += b.z;
|
||||
value.w += b.w;
|
||||
}
|
||||
dst[index] = value;
|
||||
}
|
||||
|
||||
static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
|
||||
if (!tensor)
|
||||
return true;
|
||||
size_t offset = 0;
|
||||
while (tensor->view_src) {
|
||||
offset += tensor->view_offs;
|
||||
tensor = tensor->view_src;
|
||||
}
|
||||
return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
|
||||
}
|
||||
static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
|
||||
SDI8Operation op;
|
||||
if (!sd_ggml_i8_describe(node, &op) || !aligned(node))
|
||||
return false;
|
||||
const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node);
|
||||
if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256))))
|
||||
return false;
|
||||
for (auto* input : node->src)
|
||||
if (!aligned(input))
|
||||
return false;
|
||||
const auto info = cuda_device_info(device);
|
||||
return info.int8_blas &&
|
||||
ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
|
||||
}
|
||||
static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
|
||||
int group;
|
||||
if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
|
||||
return false;
|
||||
const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
|
||||
return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0;
|
||||
}
|
||||
enum Kind { MAT,
|
||||
QUANTIZE,
|
||||
SWIGLU,
|
||||
SWIGLU_QUANTIZE,
|
||||
CONCAT };
|
||||
static bool fusion_disabled() {
|
||||
static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
|
||||
return disabled;
|
||||
}
|
||||
static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
|
||||
auto* g = const_cast<ggml_cgraph*>(graph);
|
||||
if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
|
||||
return 0;
|
||||
const auto* mat = ggml_graph_node(g, first);
|
||||
SDI8Operation op;
|
||||
if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias)
|
||||
return 0;
|
||||
const auto* next = ggml_graph_node(g, first + 1);
|
||||
if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 &&
|
||||
next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 &&
|
||||
next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0]))
|
||||
return 2;
|
||||
if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 ||
|
||||
next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) ||
|
||||
ggml_nelements(next) > int64_t(INT_MAX) * 256)
|
||||
return 0;
|
||||
if (first + 2 < ggml_graph_n_nodes(g)) {
|
||||
const auto* quant = ggml_graph_node(g, first + 2);
|
||||
if (quant->src[0] == next && convrot_supports(device, quant, nullptr) &&
|
||||
size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
|
||||
return 3;
|
||||
}
|
||||
return 2;
|
||||
}
|
||||
struct Plan {
|
||||
ggml_ext_gpu_device_info device;
|
||||
Kind kind;
|
||||
int64_t k, n, rows, padded;
|
||||
size_t quant, scales, accum, partial, output, gemm_offset;
|
||||
SDI8CudaGemmPlan* gemm = nullptr;
|
||||
bool fast_swiglu = false;
|
||||
};
|
||||
static size_t reserve(size_t& total, size_t bytes) {
|
||||
const size_t offset = total;
|
||||
total += (bytes + 255) / 256 * 256;
|
||||
return offset;
|
||||
}
|
||||
static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
|
||||
const auto* dst = nodes[count - 1];
|
||||
SDI8Operation op;
|
||||
const bool mat = sd_ggml_i8_describe(nodes[0], &op);
|
||||
auto* p = new Plan{};
|
||||
*state = p;
|
||||
auto device = ggml_backend_get_device(backend);
|
||||
p->device = cuda_device_info(device);
|
||||
p->kind = !mat ? QUANTIZE : count == 3 ? SWIGLU_QUANTIZE
|
||||
: count == 1 ? MAT
|
||||
: dst->op == GGML_OP_CONCAT ? CONCAT
|
||||
: SWIGLU;
|
||||
p->k = mat ? op.weight->ne[0] : dst->ne[0];
|
||||
p->n = mat ? nodes[0]->ne[0] : 0;
|
||||
p->rows = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]);
|
||||
p->padded = (p->rows + 3) / 4 * 4;
|
||||
size_t total = 0;
|
||||
if (mat) {
|
||||
p->quant = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0);
|
||||
p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0);
|
||||
#ifdef SD_CUDA_I8_CUTLASS
|
||||
size_t gemm_bytes = 0;
|
||||
// Retain cuBLAS when I32 accumulation can overflow.
|
||||
if (op.weight_scale && !op.bias && p->k <= INT_MAX / 16384 &&
|
||||
aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
|
||||
(op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) &&
|
||||
(p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
|
||||
if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
|
||||
p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
|
||||
p->fast_swiglu = p->gemm != nullptr;
|
||||
}
|
||||
if (!p->gemm)
|
||||
p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes);
|
||||
}
|
||||
p->gemm_offset = reserve(total, gemm_bytes);
|
||||
#endif
|
||||
const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu);
|
||||
p->accum = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
|
||||
}
|
||||
const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
|
||||
p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
|
||||
p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
|
||||
*workspace = total;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void destroy(void* state, void*) {
|
||||
auto* p = static_cast<Plan*>(state);
|
||||
#ifdef SD_CUDA_I8_CUTLASS
|
||||
if (p)
|
||||
sd_cuda_i8_gemm_free(p->gemm);
|
||||
#endif
|
||||
delete p;
|
||||
}
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
if (!b)
|
||||
return false;
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
|
||||
if (p.padded > p.rows) {
|
||||
auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
|
||||
if (error != cuda_success)
|
||||
return error;
|
||||
}
|
||||
const size_t shared = size_t(p.k) * sizeof(float);
|
||||
if (shared + 128 <= p.device.shared_memory) {
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<int, bool> configured;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
if (!configured[p.device.device]) {
|
||||
auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
|
||||
if (error != cuda_success)
|
||||
return error;
|
||||
configured[p.device.device] = true;
|
||||
}
|
||||
}
|
||||
int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256
|
||||
: 1024;
|
||||
quantize_convrot<<<unsigned(p.rows), threads, shared, stream>>>(src, dst, scales, p.k, p.rows, p.k);
|
||||
} else {
|
||||
const int64_t groups = p.k / 256, total_groups = p.rows * groups;
|
||||
convrot_group_amax_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, partial, p.k, groups, total_groups);
|
||||
reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
|
||||
convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
|
||||
}
|
||||
return cuda_last_error();
|
||||
}
|
||||
template <bool Dequantize>
|
||||
static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
|
||||
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int64_t half = n / 2;
|
||||
if (i >= rows * half)
|
||||
return;
|
||||
const int64_t row = i / half, col = i % half;
|
||||
float gate, up;
|
||||
if constexpr (Dequantize) {
|
||||
gate = float(reinterpret_cast<const int32_t*>(src)[row * n + col]) * xs[row];
|
||||
up = float(reinterpret_cast<const int32_t*>(src)[row * n + half + col]) * xs[row];
|
||||
gate *= ws[col];
|
||||
up *= ws[half + col];
|
||||
} else {
|
||||
gate = src[row * n + col];
|
||||
up = src[row * n + half + col];
|
||||
}
|
||||
dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
|
||||
}
|
||||
template <bool Swiglu, bool Dequantize, bool Rotated>
|
||||
static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
|
||||
const int64_t hidden = p.n / 2;
|
||||
auto kernel = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<int, bool> initialized;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
if (!initialized[p.device.device]) {
|
||||
auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
|
||||
if (error != cuda_success)
|
||||
return error;
|
||||
initialized[p.device.device] = true;
|
||||
}
|
||||
}
|
||||
if (p.padded > p.rows) {
|
||||
auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
|
||||
if (error != cuda_success)
|
||||
return error;
|
||||
}
|
||||
const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
|
||||
: 1024)
|
||||
: 256;
|
||||
kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
|
||||
reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
|
||||
Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
|
||||
return cuda_last_error();
|
||||
}
|
||||
static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
|
||||
const auto& p = *static_cast<const Plan*>(state);
|
||||
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<CUDAStream>(context.stream);
|
||||
auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
||||
auto* dst = nodes[count - 1];
|
||||
SDI8Operation op;
|
||||
const bool mat = sd_ggml_i8_describe(nodes[0], &op);
|
||||
bool alias = false;
|
||||
for (int i = 0; i < count; ++i)
|
||||
for (auto* input : nodes[i]->src) {
|
||||
bool internal = false;
|
||||
for (int j = 0; j < count; ++j)
|
||||
internal |= input == nodes[j];
|
||||
if (!internal)
|
||||
alias |= overlaps(dst, input);
|
||||
}
|
||||
// CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes.
|
||||
if (p.gemm && p.kind == SWIGLU_QUANTIZE)
|
||||
alias = false;
|
||||
void* output = alias ? scratch + p.output : dst->data;
|
||||
CUDAError error = cuda_success;
|
||||
if (!mat) {
|
||||
auto* quant = static_cast<int8_t*>(output);
|
||||
auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
|
||||
error = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch + p.partial));
|
||||
} else {
|
||||
auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
|
||||
const float* scales;
|
||||
if (op.input->type == GGML_TYPE_F32) {
|
||||
auto* temporary_scales = reinterpret_cast<float*>(scratch + p.scales);
|
||||
scales = temporary_scales;
|
||||
if (op.group_size)
|
||||
error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
|
||||
else {
|
||||
if (p.padded > p.rows)
|
||||
error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
|
||||
if (error == cuda_success)
|
||||
quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
|
||||
}
|
||||
} else {
|
||||
auto* data = static_cast<int8_t*>(op.input->data);
|
||||
scales = op.packed ? reinterpret_cast<const float*>(data + p.k * p.padded) : static_cast<const float*>(op.input_scale->data);
|
||||
if (op.packed || p.rows == p.padded)
|
||||
quant = data;
|
||||
else {
|
||||
error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
|
||||
if (error == cuda_success)
|
||||
error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
|
||||
}
|
||||
}
|
||||
if (error != cuda_success || cuda_last_error() != cuda_success)
|
||||
return GGML_STATUS_FAILED;
|
||||
const auto* ws = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
|
||||
const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
|
||||
const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
|
||||
auto* gemm_output = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
|
||||
auto* accum = reinterpret_cast<int32_t*>(scratch + p.accum);
|
||||
const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
|
||||
#ifdef SD_CUDA_I8_CUTLASS
|
||||
if (p.gemm) {
|
||||
auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
|
||||
intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch + p.gemm_offset, stream);
|
||||
if (status != GGML_STATUS_SUCCESS)
|
||||
return status;
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
auto blas_query = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
|
||||
auto handle = blas_query ? blas_query(execution->backend) : nullptr;
|
||||
if (!handle)
|
||||
return GGML_STATUS_FAILED;
|
||||
#if defined(SD_CUDA_VENDOR_MUSA)
|
||||
return GGML_STATUS_FAILED;
|
||||
#else
|
||||
const int32_t alpha = 1, beta = 0;
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
#if HIP_VERSION >= 60500000
|
||||
constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
|
||||
#else
|
||||
constexpr auto compute_type = HIPBLAS_R_32I;
|
||||
#endif
|
||||
const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
|
||||
int(p.n), int(p.padded), int(p.k), &alpha,
|
||||
op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
|
||||
&beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
|
||||
if (status != HIPBLAS_STATUS_SUCCESS)
|
||||
#else
|
||||
const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
int(p.n), int(p.padded), int(p.k), &alpha,
|
||||
op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
|
||||
&beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
|
||||
if (status != CUBLAS_STATUS_SUCCESS)
|
||||
#endif
|
||||
return GGML_STATUS_FAILED;
|
||||
#endif
|
||||
}
|
||||
const auto* values = reinterpret_cast<const float*>(accum);
|
||||
if (p.kind == SWIGLU_QUANTIZE) {
|
||||
if (p.fast_swiglu)
|
||||
error = quantize_epilogue<false, false, true>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
|
||||
else if (p.gemm)
|
||||
error = quantize_epilogue<true, false, false>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
|
||||
else
|
||||
error = quantize_epilogue<true, true, false>(p, values, scales, ws, static_cast<int8_t*>(output), stream);
|
||||
} else if (p.kind == SWIGLU && !p.fast_swiglu) {
|
||||
const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256);
|
||||
if (p.gemm)
|
||||
swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
|
||||
else
|
||||
swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
|
||||
error = cuda_last_error();
|
||||
} else if (!p.gemm) {
|
||||
const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
|
||||
if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
|
||||
dequantize_i32_rows_vec4_cuda<<<dim3(unsigned((p.n + 1023) / 1024), gy, gz), 256, 0, stream>>>(
|
||||
reinterpret_cast<const int4*>(accum), scales, reinterpret_cast<const float4*>(ws), reinterpret_cast<const float4*>(bias), reinterpret_cast<float4*>(gemm_output), p.n / 4, p.rows);
|
||||
} else {
|
||||
dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
|
||||
}
|
||||
error = cuda_last_error();
|
||||
}
|
||||
if (error == cuda_success && prefix)
|
||||
error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
|
||||
}
|
||||
if (error == cuda_success && alias)
|
||||
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
|
||||
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
void register_i8(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.linear_i8." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_i8_op();
|
||||
desc.priority = 100;
|
||||
desc.supports = sd_i8_cuda::supports;
|
||||
desc.prepare = sd_i8_cuda::prepare;
|
||||
desc.execute = sd_i8_cuda::execute;
|
||||
desc.destroy = sd_i8_cuda::destroy;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
|
||||
desc.priority = 240;
|
||||
desc.supports = nullptr;
|
||||
desc.match = sd_i8_cuda::match;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.priority = 100;
|
||||
desc.match = nullptr;
|
||||
desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
|
||||
desc.custom_op = sd_ggml_convrot_op();
|
||||
desc.supports = sd_i8_cuda::convrot_supports;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
#include "runtime.h"
|
||||
|
||||
#ifdef SD_CUDA_CUTE_ATTENTION
|
||||
void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
|
||||
void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
|
||||
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
|
||||
#endif
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
void register_rope(ggml_backend_dev_t device);
|
||||
void register_i8(ggml_backend_dev_t device);
|
||||
void register_convrot(ggml_backend_dev_t device);
|
||||
void register_affine(ggml_backend_dev_t device);
|
||||
#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
|
||||
void register_sage_attention(ggml_backend_dev_t device);
|
||||
void register_sol_attention(ggml_backend_dev_t device);
|
||||
#endif
|
||||
|
||||
void register_extensions(ggml_backend_dev_t device) {
|
||||
if (cuda_device_info(device).device < 0)
|
||||
return;
|
||||
register_affine(device);
|
||||
register_i8(device);
|
||||
register_convrot(device);
|
||||
register_rope(device);
|
||||
#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
|
||||
register_sage_attention(device);
|
||||
register_sol_attention(device);
|
||||
#endif
|
||||
#ifdef SD_CUDA_CUTE_ATTENTION
|
||||
sd_ggml_register_cuda_i8_projection(device);
|
||||
sd_ggml_register_cuda_i8_gated_residual(device);
|
||||
sd_ggml_register_cuda_attention(device);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
#ifndef __SD_CORE_CUDA_REGISTRATION_H__
|
||||
#define __SD_CORE_CUDA_REGISTRATION_H__
|
||||
#include "ggml-backend.h"
|
||||
namespace sd_cuda {
|
||||
void register_extensions(ggml_backend_dev_t device);
|
||||
}
|
||||
namespace sd_hip {
|
||||
void register_extensions(ggml_backend_dev_t device);
|
||||
}
|
||||
namespace sd_musa {
|
||||
void register_extensions(ggml_backend_dev_t device);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,337 @@
|
||||
#include "core/cuda/rope.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <type_traits>
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
static __device__ __forceinline__ float sd_warp_sum(float v) {
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
|
||||
return v;
|
||||
}
|
||||
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 = cuda_shuffle_xor(value, 1, 32);
|
||||
if (col < 2 * rot_half) {
|
||||
const float* m = pe + 4 * (token * rot_half + col / 2);
|
||||
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct RopeApplyStrides {
|
||||
int64_t x[4];
|
||||
int64_t pe[4];
|
||||
};
|
||||
|
||||
template <bool interleaved>
|
||||
static __global__ void rope_apply_f32(
|
||||
const float* x,
|
||||
const float* pe,
|
||||
float* dst,
|
||||
int64_t half_dim,
|
||||
int64_t rot_half,
|
||||
int64_t heads,
|
||||
int64_t tokens,
|
||||
int64_t pairs,
|
||||
RopeApplyStrides strides) {
|
||||
const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (index >= pairs) {
|
||||
return;
|
||||
}
|
||||
const int64_t j = index % half_dim;
|
||||
const int64_t row = index / half_dim;
|
||||
const int64_t token = row % tokens;
|
||||
const int64_t head = (row / tokens) % heads;
|
||||
const int64_t batch = row / (tokens * heads);
|
||||
const int64_t i0 = interleaved || j >= rot_half ? 2 * j : j;
|
||||
const int64_t i1 = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
|
||||
const float* x_row = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
|
||||
const float a = x_row[i0 * strides.x[0]];
|
||||
const float b = x_row[i1 * strides.x[0]];
|
||||
float* out = dst + row * half_dim * 2;
|
||||
if (j >= rot_half) {
|
||||
out[i0] = a;
|
||||
out[i1] = b;
|
||||
return;
|
||||
}
|
||||
const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
|
||||
const float c = matrix[0];
|
||||
const float s = matrix[strides.pe[1]];
|
||||
// Match the separate mul/add rounding of the unfused graph.
|
||||
if (interleaved) {
|
||||
out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
|
||||
out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
|
||||
} else {
|
||||
out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
|
||||
out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
|
||||
}
|
||||
}
|
||||
|
||||
static bool rope_interleaved(const ggml_tensor* rope) {
|
||||
int32_t value;
|
||||
std::memcpy(&value, ggml_ext_op_attributes(rope), sizeof(value));
|
||||
return value != 0;
|
||||
}
|
||||
|
||||
static float param_f32(const ggml_tensor* node, int index) {
|
||||
float value;
|
||||
std::memcpy(&value, node->op_params + index, sizeof(value));
|
||||
return value;
|
||||
}
|
||||
|
||||
static bool rope_is_node(const ggml_tensor* node) {
|
||||
return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
|
||||
}
|
||||
|
||||
static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
|
||||
const auto* x = rope->src[0];
|
||||
const auto* pe = rope->src[1];
|
||||
if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
|
||||
x->ne[0] % 2 || 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] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
|
||||
rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
|
||||
(ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
|
||||
return false;
|
||||
}
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool rope_fusion_supports(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_is_node(rope) ||
|
||||
mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
|
||||
return false;
|
||||
const auto* x = norm->src[0];
|
||||
const auto* weight = mul->src[1];
|
||||
return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
|
||||
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
|
||||
ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
|
||||
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
|
||||
x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
|
||||
x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
|
||||
(ggml_nrows(x) + 3) / 4 <= INT_MAX;
|
||||
}
|
||||
|
||||
template <typename T, bool Interleaved>
|
||||
static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
|
||||
const auto* x = nodes[0]->src[0];
|
||||
const auto* pe = nodes[2]->src[1];
|
||||
const int64_t rows = ggml_nrows(x);
|
||||
#define SD_ROPE(D) \
|
||||
case D: \
|
||||
rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
|
||||
static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data), \
|
||||
static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows, \
|
||||
x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
|
||||
break
|
||||
switch (x->ne[0]) {
|
||||
SD_ROPE(32);
|
||||
SD_ROPE(64);
|
||||
SD_ROPE(96);
|
||||
SD_ROPE(128);
|
||||
SD_ROPE(160);
|
||||
SD_ROPE(192);
|
||||
SD_ROPE(224);
|
||||
SD_ROPE(256);
|
||||
default:
|
||||
return cuda_invalid_value;
|
||||
}
|
||||
#undef SD_ROPE
|
||||
return cuda_last_error();
|
||||
}
|
||||
|
||||
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
|
||||
if (fp16) {
|
||||
return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
|
||||
: rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
|
||||
}
|
||||
return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
|
||||
: rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
|
||||
}
|
||||
|
||||
static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
|
||||
static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
|
||||
if (disabled)
|
||||
return 0;
|
||||
auto* g = const_cast<ggml_cgraph*>(graph);
|
||||
if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
|
||||
ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
|
||||
return 0;
|
||||
if (first + 4 < ggml_graph_n_nodes(g)) {
|
||||
const auto* rope = ggml_graph_node(g, first + 2);
|
||||
const auto* scale = ggml_graph_node(g, first + 3);
|
||||
const auto* copy = ggml_graph_node(g, first + 4);
|
||||
if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
|
||||
scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
|
||||
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
|
||||
return 5;
|
||||
}
|
||||
return 3;
|
||||
}
|
||||
|
||||
static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
|
||||
static const bool disabled = [] {
|
||||
const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
|
||||
return value && std::atoi(value) != 0;
|
||||
}();
|
||||
if (disabled)
|
||||
return;
|
||||
const int total = ggml_graph_n_nodes(graph);
|
||||
for (int i = 2; i + 1 < total; ++i) {
|
||||
const auto* mul = ggml_graph_node(graph, i);
|
||||
const auto* rope = ggml_graph_node(graph, i + 1);
|
||||
if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
|
||||
!mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
|
||||
continue;
|
||||
const auto* norm = mul->src[0];
|
||||
for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
|
||||
if (ggml_graph_node(graph, j) != norm)
|
||||
continue;
|
||||
if (j + 1 < i) {
|
||||
if (j >= 2) {
|
||||
const auto* mat = ggml_graph_node(graph, j - 2);
|
||||
const auto* view = ggml_graph_node(graph, j - 1);
|
||||
if (ggml_ext_op_get(mat) == sd_ggml_i8_op() &&
|
||||
view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
|
||||
ggml_ext_graph_delay(graph, j - 2, 3, i))
|
||||
break;
|
||||
}
|
||||
ggml_ext_graph_delay(graph, j, 1, i);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
|
||||
*workspace = ggml_nbytes(nodes[count - 1]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
static void rope_destroy(void*, void*) {}
|
||||
|
||||
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
|
||||
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
|
||||
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
|
||||
}
|
||||
|
||||
static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto stream = static_cast<CUDAStream>(context.stream);
|
||||
const auto* rope = nodes[count == 1 ? 0 : 2];
|
||||
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
|
||||
auto* dst = nodes[count - 1];
|
||||
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
|
||||
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
|
||||
CUDAError error;
|
||||
if (count > 1) {
|
||||
error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
|
||||
} else {
|
||||
RopeApplyStrides strides;
|
||||
const auto* pe = rope->src[1];
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
strides.x[i] = x->nb[i] / 4;
|
||||
strides.pe[i] = pe->nb[i] / 4;
|
||||
}
|
||||
const int64_t pairs = ggml_nelements(x) / 2;
|
||||
auto kernel = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
|
||||
kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
|
||||
static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
|
||||
error = cuda_last_error();
|
||||
}
|
||||
if (error == cuda_success && alias)
|
||||
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
|
||||
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
|
||||
void register_rope(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.rope_apply." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_rope_op();
|
||||
desc.priority = 100;
|
||||
desc.supports = rope_supports;
|
||||
desc.prepare = rope_prepare;
|
||||
desc.destroy = rope_destroy;
|
||||
desc.execute = rope_execute;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
desc.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
|
||||
desc.op = GGML_OP_RMS_NORM;
|
||||
desc.custom_op = nullptr;
|
||||
desc.priority = 250;
|
||||
desc.supports = nullptr;
|
||||
desc.match = rope_match;
|
||||
desc.optimize = rope_optimize;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
#ifndef __SD_CORE_CUDA_ROPE_H__
|
||||
#define __SD_CORE_CUDA_ROPE_H__
|
||||
|
||||
#include "runtime.h"
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
|
||||
CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
|
||||
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,111 @@
|
||||
#ifndef __SD_CORE_CUDA_RUNTIME_H__
|
||||
#define __SD_CORE_CUDA_RUNTIME_H__
|
||||
|
||||
#include "ggml-ext.h"
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hip/hip_runtime.h>
|
||||
#define SD_CUDA_NAMESPACE sd_hip
|
||||
#define SD_CUDA_NAME "hip"
|
||||
#elif defined(SD_CUDA_VENDOR_MUSA)
|
||||
#include <musa_fp16.h>
|
||||
#include <musa_runtime.h>
|
||||
#define SD_CUDA_NAMESPACE sd_musa
|
||||
#define SD_CUDA_NAME "musa"
|
||||
#else
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#define SD_CUDA_NAMESPACE sd_cuda
|
||||
#define SD_CUDA_NAME "cuda"
|
||||
#endif
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
using CUDAError = hipError_t;
|
||||
using CUDAStream = hipStream_t;
|
||||
static constexpr auto cuda_success = hipSuccess;
|
||||
static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
|
||||
static constexpr auto cuda_kind = GGML_EXT_GPU_HIP;
|
||||
static CUDAError cuda_last_error() {
|
||||
return hipGetLastError();
|
||||
}
|
||||
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
|
||||
return hipMemsetAsync(dst, value, bytes, stream);
|
||||
}
|
||||
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
|
||||
return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
|
||||
}
|
||||
template <typename Kernel>
|
||||
static CUDAError cuda_set_shared_memory(Kernel, int) {
|
||||
// HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
|
||||
return cuda_success;
|
||||
}
|
||||
#elif defined(SD_CUDA_VENDOR_MUSA)
|
||||
using CUDAError = musaError_t;
|
||||
using CUDAStream = musaStream_t;
|
||||
static constexpr auto cuda_success = musaSuccess;
|
||||
static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
|
||||
static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA;
|
||||
static CUDAError cuda_last_error() {
|
||||
return musaGetLastError();
|
||||
}
|
||||
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
|
||||
return musaMemsetAsync(dst, value, bytes, stream);
|
||||
}
|
||||
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
|
||||
return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
|
||||
}
|
||||
template <typename Kernel>
|
||||
static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
|
||||
return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
|
||||
}
|
||||
#else
|
||||
using CUDAError = cudaError_t;
|
||||
using CUDAStream = cudaStream_t;
|
||||
static constexpr auto cuda_success = cudaSuccess;
|
||||
static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
|
||||
static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA;
|
||||
static CUDAError cuda_last_error() {
|
||||
return cudaGetLastError();
|
||||
}
|
||||
static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
|
||||
return cudaMemsetAsync(dst, value, bytes, stream);
|
||||
}
|
||||
static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
|
||||
return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
|
||||
}
|
||||
template <typename Kernel>
|
||||
static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
|
||||
return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
|
||||
}
|
||||
#endif
|
||||
template <typename T>
|
||||
static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
return __shfl(value, lane, width);
|
||||
#else
|
||||
return __shfl_sync(0xffffffff, value, lane, width);
|
||||
#endif
|
||||
}
|
||||
template <typename T>
|
||||
static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
return __shfl_xor(value, lane, width);
|
||||
#else
|
||||
return __shfl_xor_sync(0xffffffff, value, lane, width);
|
||||
#endif
|
||||
}
|
||||
static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
|
||||
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
|
||||
auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
|
||||
return query && query(backend, context) && context->kind == cuda_kind;
|
||||
}
|
||||
static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
|
||||
ggml_ext_gpu_device_info info = {};
|
||||
auto query = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
|
||||
if (!query || !query(device, &info) || info.kind != cuda_kind)
|
||||
return {cuda_kind, -1, 0, 0, false};
|
||||
return info;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,387 @@
|
||||
#include "core/ggml_extensions.h"
|
||||
#include "runtime.h"
|
||||
|
||||
#if __CUDACC_VER_MAJOR__ >= 12
|
||||
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
#include <cstring>
|
||||
#include <map>
|
||||
#include "sageattention/attn.cuh"
|
||||
#include "sageattention/attn_fp8.cuh"
|
||||
|
||||
static constexpr int sage_cuda_version = __CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10;
|
||||
|
||||
namespace sd_cuda {
|
||||
|
||||
struct SageDevice {
|
||||
int cc;
|
||||
int compiled_cc;
|
||||
};
|
||||
|
||||
static ggml_sage_attn_mode sage_mode(const SageDevice& device, ggml_sage_attn_mode requested) {
|
||||
if (requested != GGML_SAGE_ATTN_AUTO)
|
||||
return requested;
|
||||
if (sage_cuda_version >= 12040 && device.cc >= 89 && device.compiled_cc >= 89) {
|
||||
if (sage_cuda_version >= 12080 && device.cc != 90)
|
||||
return GGML_SAGE_ATTN_2_PLUS_PLUS;
|
||||
return GGML_SAGE_ATTN_2;
|
||||
}
|
||||
return GGML_SAGE_ATTN_AUTO;
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ float warp_reduce_max(float x) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset));
|
||||
return x;
|
||||
}
|
||||
|
||||
static __global__ void sage_output(const half* input, float* output, int count) {
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < count)
|
||||
output[i] = __half2float(input[i]);
|
||||
}
|
||||
|
||||
struct SageWorkspace {
|
||||
size_t q;
|
||||
size_t k;
|
||||
size_t q_scale;
|
||||
size_t k_scale;
|
||||
size_t k_mean;
|
||||
size_t v;
|
||||
size_t v_scale;
|
||||
size_t output;
|
||||
size_t size;
|
||||
|
||||
explicit SageWorkspace(const ggml_tensor* dst) {
|
||||
const ggml_tensor* qt = dst->src[0];
|
||||
const ggml_tensor* kt = dst->src[1];
|
||||
size = 0;
|
||||
q = reserve(ggml_nelements(qt));
|
||||
k = reserve(ggml_nelements(kt));
|
||||
q_scale = reserve(GGML_PAD(qt->ne[1], 128) / 32 * 8 * qt->ne[2] * qt->ne[3] * sizeof(float));
|
||||
k_scale = reserve(GGML_PAD(kt->ne[1], 64) / 64 * 4 * kt->ne[2] * kt->ne[3] * sizeof(float));
|
||||
k_mean = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
|
||||
v = reserve(kt->ne[0] * GGML_PAD(kt->ne[1], 64) * kt->ne[2] * kt->ne[3]);
|
||||
v_scale = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
|
||||
output = reserve(ggml_nelements(dst) * sizeof(half));
|
||||
}
|
||||
|
||||
size_t reserve(size_t bytes) {
|
||||
const size_t offset = size;
|
||||
size += GGML_PAD(bytes, 256);
|
||||
return offset;
|
||||
}
|
||||
};
|
||||
|
||||
template <int D>
|
||||
static __global__ void sage_key_mean(const float* k, float* mean, int length) {
|
||||
const int channel = blockIdx.x * 32 + threadIdx.x % 32;
|
||||
const int row = threadIdx.x / 32;
|
||||
const size_t base = size_t(blockIdx.y) * length * D;
|
||||
float sum = 0.0f;
|
||||
for (int token = row; token < length; token += 8) {
|
||||
sum += k[base + size_t(token) * D + channel];
|
||||
}
|
||||
__shared__ float sums[256];
|
||||
sums[threadIdx.x] = sum;
|
||||
__syncthreads();
|
||||
if (row == 0) {
|
||||
for (int i = 1; i < 8; ++i) {
|
||||
sum += sums[threadIdx.x + i * 32];
|
||||
}
|
||||
mean[size_t(blockIdx.y) * D + channel] = sum / length;
|
||||
}
|
||||
}
|
||||
|
||||
template <int D, int TOKENS, bool SMOOTH>
|
||||
static __global__ void sage_quantize(const float* input, const float* mean, int8_t* output, float* scales, int length) {
|
||||
constexpr int VALUES = D * TOKENS / 256;
|
||||
const int head = blockIdx.y;
|
||||
const int first = blockIdx.x * TOKENS * D;
|
||||
const size_t base = size_t(head) * length * D;
|
||||
float values[VALUES];
|
||||
float amax = 1e-7f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VALUES; ++i) {
|
||||
const int index = first + threadIdx.x + i * 256;
|
||||
float value = 0.0f;
|
||||
if (index < length * D) {
|
||||
value = input[base + index];
|
||||
if constexpr (SMOOTH) {
|
||||
value -= mean[size_t(head) * D + index % D];
|
||||
}
|
||||
}
|
||||
values[i] = value;
|
||||
amax = fmaxf(amax, fabsf(value));
|
||||
}
|
||||
amax = warp_reduce_max(amax);
|
||||
__shared__ float maxima[8];
|
||||
if (threadIdx.x % 32 == 0) {
|
||||
maxima[threadIdx.x / 32] = amax;
|
||||
}
|
||||
__syncthreads();
|
||||
if (threadIdx.x < 32) {
|
||||
amax = threadIdx.x < 8 ? maxima[threadIdx.x] : 1e-7f;
|
||||
amax = warp_reduce_max(amax);
|
||||
if (threadIdx.x == 0) {
|
||||
maxima[0] = amax;
|
||||
scales[size_t(head) * gridDim.x + blockIdx.x] = amax / 127.0f;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const float inverse_scale = 127.0f / maxima[0];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VALUES; ++i) {
|
||||
const int index = first + threadIdx.x + i * 256;
|
||||
if (index < length * D) {
|
||||
output[base + index] = int8_t(__float2int_rn(values[i] * inverse_scale));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int D, bool KEY>
|
||||
static __global__ void sage_quantize_per_thread(const float* input, const float* mean, int8_t* output, float* scales, int length) {
|
||||
constexpr int GROUPS = KEY ? 4 : 8;
|
||||
constexpr int TOKENS = KEY ? 64 : 32;
|
||||
constexpr int VALUES = D * TOKENS / GROUPS / 32;
|
||||
const int group = threadIdx.x / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
const size_t base = size_t(blockIdx.y) * length * D;
|
||||
float values[VALUES];
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VALUES; ++i) {
|
||||
const int index = lane + i * 32;
|
||||
const int row = index / D;
|
||||
const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
|
||||
float value = token < length ? input[base + size_t(token) * D + index % D] : 0.0f;
|
||||
if constexpr (KEY) {
|
||||
if (token < length) {
|
||||
value -= mean[size_t(blockIdx.y) * D + index % D];
|
||||
}
|
||||
}
|
||||
values[i] = value;
|
||||
amax = fmaxf(amax, fabsf(value));
|
||||
}
|
||||
const float scale = warp_reduce_max(amax) / 127.0f + 1e-7f;
|
||||
if (lane == 0) {
|
||||
scales[(size_t(blockIdx.y) * gridDim.x + blockIdx.x) * GROUPS + group] = scale;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VALUES; ++i) {
|
||||
const int index = lane + i * 32;
|
||||
const int row = index / D;
|
||||
const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
|
||||
if (token < length) {
|
||||
const float value = values[i] / scale;
|
||||
output[base + size_t(token) * D + index % D] = int8_t(__float2int_rz(value + copysignf(0.5f, value)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int D>
|
||||
static __global__ void sage_value_scale(const half* v, float* scales, int length, float scale_max) {
|
||||
const int channel = blockIdx.x * 32 + threadIdx.x % 32;
|
||||
const int row = threadIdx.x / 32;
|
||||
const size_t base = size_t(blockIdx.y) * length * D;
|
||||
float amax = 1e-7f;
|
||||
for (int token = row; token < length; token += 8) {
|
||||
amax = fmaxf(amax, fabsf(__half2float(v[base + size_t(token) * D + channel])));
|
||||
}
|
||||
__shared__ float maxima[256];
|
||||
maxima[threadIdx.x] = amax;
|
||||
__syncthreads();
|
||||
if (row == 0) {
|
||||
for (int i = 1; i < 8; ++i) {
|
||||
amax = fmaxf(amax, maxima[threadIdx.x + i * 32]);
|
||||
}
|
||||
scales[size_t(blockIdx.y) * D + channel] = amax / scale_max;
|
||||
}
|
||||
}
|
||||
|
||||
template <int D>
|
||||
static __global__ void sage_quantize_value(const half* v, const float* scales, int8_t* output, int length, int padded) {
|
||||
__shared__ uint8_t tile[32][D + 4];
|
||||
const int head = blockIdx.y;
|
||||
const int first = blockIdx.x * 32;
|
||||
const size_t base = size_t(head) * length * D;
|
||||
for (int i = threadIdx.x; i < 32 * D; i += 256) {
|
||||
const int token = first + i / D;
|
||||
const int channel = i % D;
|
||||
const float value = token < length ? __half2float(v[base + size_t(token) * D + channel]) / scales[size_t(head) * D + channel] : 0.0f;
|
||||
// FP8 MMA consumes pairs of tokens in a permuted, transposed V tile.
|
||||
const int row = i / D;
|
||||
const int permuted = row / 16 * 16 + row % 16 / 8 * 2 + row % 8 / 2 * 4 + row % 2;
|
||||
tile[permuted][channel] = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
|
||||
}
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < 32 * D; i += 256) {
|
||||
output[(size_t(head) * D + i / 32) * padded + first + i % 32] = tile[i % 32][i / 32];
|
||||
}
|
||||
}
|
||||
|
||||
template <int D, bool PLUS_PLUS>
|
||||
static void sage_attn_fp8(cudaStream_t stream, ggml_tensor* dst, char* data, const SageWorkspace& work, float scale) {
|
||||
const ggml_tensor* q = dst->src[0];
|
||||
const ggml_tensor* k = dst->src[1];
|
||||
const ggml_tensor* v = dst->src[2];
|
||||
const int lq = q->ne[1], lk = k->ne[1], padded = GGML_PAD(lk, 64);
|
||||
const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
|
||||
auto* qi = reinterpret_cast<int8_t*>(data + work.q);
|
||||
auto* ki = reinterpret_cast<int8_t*>(data + work.k);
|
||||
auto* vi = reinterpret_cast<int8_t*>(data + work.v);
|
||||
auto* qs = reinterpret_cast<float*>(data + work.q_scale);
|
||||
auto* ks = reinterpret_cast<float*>(data + work.k_scale);
|
||||
auto* vs = reinterpret_cast<float*>(data + work.v_scale);
|
||||
auto* km = reinterpret_cast<float*>(data + work.k_mean);
|
||||
auto* out = reinterpret_cast<half*>(data + work.output);
|
||||
sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
|
||||
sage_quantize_per_thread<D, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
|
||||
static_cast<const float*>(q->data), nullptr, qi, qs, lq);
|
||||
sage_quantize_per_thread<D, true><<<dim3(padded / 64, hk * batch), 128, 0, stream>>>(
|
||||
static_cast<const float*>(k->data), km, ki, ks, lk);
|
||||
// The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++.
|
||||
sage_value_scale<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f);
|
||||
sage_quantize_value<D><<<dim3(padded / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, vi, lk, padded);
|
||||
auto kernel = sageattention::qk_int_sv_f8_attn_kernel<128, 64, 32, 64, D,
|
||||
sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerThread,
|
||||
sageattention::QuantGranularity::kPerThread, float, true, half,
|
||||
sageattention::ComputeUnit::kCudaCore, sageattention::MaskMode::kNone, false, true, false, PLUS_PLUS>;
|
||||
kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), 256 * D, stream>>>(
|
||||
qi, ki, vi, out, nullptr, qs, ks, vs, nullptr, lq, lk, hq / hk,
|
||||
D * lq * hq, D, D * lq, D * lk * hk, D, D * lk,
|
||||
D * padded * hk, D * padded, padded, D * lq * hq, D * hq, D, scale);
|
||||
sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
|
||||
}
|
||||
|
||||
template <int D>
|
||||
static void sage_attn(cudaStream_t stream, ggml_tensor* dst, char* data, ggml_sage_attn_mode mode, float scale) {
|
||||
const ggml_tensor* q = dst->src[0];
|
||||
const ggml_tensor* k = dst->src[1];
|
||||
const ggml_tensor* v = dst->src[2];
|
||||
const int lq = q->ne[1], lk = k->ne[1];
|
||||
const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
|
||||
const SageWorkspace work(dst);
|
||||
if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS) {
|
||||
sage_attn_fp8<D, true>(stream, dst, data, work, scale);
|
||||
return;
|
||||
}
|
||||
if (mode == GGML_SAGE_ATTN_2) {
|
||||
sage_attn_fp8<D, false>(stream, dst, data, work, scale);
|
||||
return;
|
||||
}
|
||||
auto* qi = reinterpret_cast<int8_t*>(data + work.q);
|
||||
auto* ki = reinterpret_cast<int8_t*>(data + work.k);
|
||||
auto* qs = reinterpret_cast<float*>(data + work.q_scale);
|
||||
auto* ks = reinterpret_cast<float*>(data + work.k_scale);
|
||||
auto* km = reinterpret_cast<float*>(data + work.k_mean);
|
||||
auto* out = reinterpret_cast<half*>(data + work.output);
|
||||
|
||||
sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
|
||||
sage_quantize<D, 32, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
|
||||
static_cast<const float*>(q->data), nullptr, qi, qs, lq);
|
||||
sage_quantize<D, 64, true><<<dim3(GGML_PAD(lk, 64) / 64, hk * batch), 256, 0, stream>>>(
|
||||
static_cast<const float*>(k->data), km, ki, ks, lk);
|
||||
|
||||
constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
|
||||
auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
|
||||
sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
|
||||
sageattention::QuantGranularity::kPerWarp, float, false, half,
|
||||
sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
|
||||
kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
|
||||
qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
|
||||
lq, lk, hq / hk,
|
||||
D * lq * hq, D, D * lq,
|
||||
D * lk * hk, D, D * lk,
|
||||
D * lk * hk, D, D * lk,
|
||||
D * lq * hq, D * hq, D,
|
||||
scale);
|
||||
sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
|
||||
}
|
||||
|
||||
static bool sage_supports(ggml_backend_dev_t, const ggml_tensor* op, void* user_data) {
|
||||
const auto& device = *static_cast<const SageDevice*>(user_data);
|
||||
SDSageAttentionParams params;
|
||||
if (!sd_ggml_sage_attn_describe(op, ¶ms) || device.cc < 80 || device.compiled_cc < 80)
|
||||
return false;
|
||||
const auto mode = sage_mode(device, params.mode);
|
||||
if (mode != GGML_SAGE_ATTN_AUTO && (sage_cuda_version < 12040 || device.cc < 89 || device.compiled_cc < 89))
|
||||
return false;
|
||||
if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS && sage_cuda_version < 12080)
|
||||
return false;
|
||||
const auto* q = op->src[0];
|
||||
const auto* k = op->src[1];
|
||||
const auto* v = op->src[2];
|
||||
if (q->ne[0] != 64 && q->ne[0] != 128)
|
||||
return false;
|
||||
for (const auto* t : {q, k, v, op}) {
|
||||
if (ggml_nelements(t) > INT_MAX || t->ne[2] * t->ne[3] > 65535 ||
|
||||
(reinterpret_cast<uintptr_t>(t->data) & 15))
|
||||
return false;
|
||||
}
|
||||
return k->ne[0] * GGML_PAD(k->ne[1], 64) * k->ne[2] * k->ne[3] <= INT_MAX;
|
||||
}
|
||||
|
||||
static ggml_status sage_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
|
||||
*workspace = SageWorkspace(nodes[0]).size;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
static void sage_destroy(void*, void*) {}
|
||||
|
||||
static ggml_status sage_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void* user_data) {
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
SDSageAttentionParams params;
|
||||
if (!sd_ggml_sage_attn_describe(nodes[0], ¶ms))
|
||||
return GGML_STATUS_FAILED;
|
||||
const auto mode = sage_mode(*static_cast<const SageDevice*>(user_data), params.mode);
|
||||
auto* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
||||
auto stream = static_cast<cudaStream_t>(context.stream);
|
||||
if (nodes[0]->src[0]->ne[0] == 64)
|
||||
sage_attn<64>(stream, nodes[0], data, mode, params.scale);
|
||||
else
|
||||
sage_attn<128>(stream, nodes[0], data, mode, params.scale);
|
||||
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
|
||||
void register_sage_attention(ggml_backend_dev_t device) {
|
||||
const auto info = cuda_device_info(device);
|
||||
if (info.device < 0)
|
||||
return;
|
||||
int previous;
|
||||
cudaDeviceProp properties;
|
||||
cudaFuncAttributes attributes;
|
||||
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
|
||||
cudaSetDevice(info.device) != cudaSuccess)
|
||||
return;
|
||||
const auto error = cudaFuncGetAttributes(&attributes, sage_output);
|
||||
if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess)
|
||||
return;
|
||||
// PTX compiled for an older target retains that target's instruction guards after JIT.
|
||||
static std::map<int, SageDevice> devices;
|
||||
auto& data = devices[info.device];
|
||||
data = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)};
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.sage_attn.cuda.v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_sage_attn_op();
|
||||
desc.priority = 100;
|
||||
desc.supports = sage_supports;
|
||||
desc.prepare = sage_prepare;
|
||||
desc.destroy = sage_destroy;
|
||||
desc.execute = sage_execute;
|
||||
desc.user_data = &data;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
|
||||
}
|
||||
#else
|
||||
namespace sd_cuda {
|
||||
void register_sage_attention(ggml_backend_dev_t) {}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright 2024 Jintao Zhang, Haofeng Huang
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,28 @@
|
||||
# SageAttention CUDA device code
|
||||
|
||||
Adapted from [thu-ml/SageAttention](https://github.com/thu-ml/SageAttention),
|
||||
revision `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`, under Apache-2.0 (see LICENSE).
|
||||
|
||||
`attn.cuh` contains the device kernel from
|
||||
`csrc/qattn/qk_int_sv_f16_cuda_sm80.cu`; `attn_fp8.cuh` adapts
|
||||
`csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh`. The remaining headers come from `csrc/`
|
||||
and `csrc/qattn/attn_utils.cuh`. PyTorch includes and host wrappers were removed;
|
||||
include paths and guards were adjusted, device helpers were placed in the
|
||||
`sageattention` namespace, and `div_ceil` was renamed to avoid macro collisions.
|
||||
Unused `<cuda/pipeline>` includes were removed so multi-architecture builds can
|
||||
also compile targets below SM70; the backend rejects SageAttention below SM80.
|
||||
Predicated loads now zero unused FP16 V lanes: masked probabilities cannot
|
||||
suppress NaNs left in shared memory by previous kernels.
|
||||
|
||||
The sd.cpp adapter in `../sage_attention.cu` supplies FP32 K smoothing and native CUDA
|
||||
per-thread INT8 quantization, matching the token grouping in upstream
|
||||
`sageattention/triton/quant_per_thread.py`. V is quantized per channel and
|
||||
transposed, padded with zeros, and permuted according to
|
||||
`csrc/fused/fused.cu`. Its FP8 range is 448 for SageAttention2 and 2.25 for
|
||||
SageAttention2++. Zero channels use a positive scale to avoid division by zero.
|
||||
SageAttention2 uses an FP32 instruction buffer; 2++ uses an FP16 instruction
|
||||
buffer added to FP32 accumulators. The SM80 compatibility path retains per-warp
|
||||
INT8 Q/K and FP16 V with FP32 accumulation.
|
||||
|
||||
The adapter requests workspace and the current stream through public ggml extension APIs; it has no Python, PyTorch,
|
||||
or Triton dependency. The public output is converted back to FP32.
|
||||
@@ -0,0 +1,592 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
|
||||
|
||||
#include "attn_utils.cuh"
|
||||
#include "math.cuh"
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
#define PACK_SIZE_QK 16 // as if it is int8
|
||||
#define PACK_SIZE_V 8 // fp16
|
||||
#define PACK_SIZE_O 8 // fp16
|
||||
|
||||
// treat as if int8 tensor core
|
||||
#define MMA_QK_M 16
|
||||
#define MMA_QK_N 16
|
||||
#define MMA_QK_K 32
|
||||
|
||||
// fp16 tensor core
|
||||
#define MMA_SV_M 16
|
||||
#define MMA_SV_N 16
|
||||
#define MMA_SV_K 16
|
||||
|
||||
template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean = false>
|
||||
__global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
|
||||
// Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes.
|
||||
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
|
||||
// compile time check
|
||||
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
|
||||
static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
|
||||
static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
|
||||
static_assert(std::is_same<DTypeSVAccum, float>::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float");
|
||||
static_assert(std::is_same<DTypeSVAccum, float>::value || std::is_same<DTypeSVAccum, half>::value, "DTypeSVAccum must be float or half");
|
||||
static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
|
||||
static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
|
||||
static_assert(!fuse_v_mean || std::is_same<DTypeSVAccum, half>::value, "fuse_v_mean only supports half");
|
||||
static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation
|
||||
|
||||
using DTypeOut2 = typename std::conditional<std::is_same<DTypeOut, half>::value, half2, nv_bfloat162>::type;
|
||||
|
||||
constexpr uint32_t num_warps_q = CTA_Q / WARP_Q;
|
||||
constexpr uint32_t num_warps_k = CTA_K / WARP_K;
|
||||
constexpr uint32_t num_warps = num_warps_q * num_warps_k;
|
||||
constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M;
|
||||
constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N;
|
||||
constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
|
||||
constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N;
|
||||
|
||||
constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
|
||||
constexpr uint32_t O_SMEM_STRIDE = head_dim;
|
||||
constexpr uint32_t V_SMEM_STRIDE = head_dim;
|
||||
|
||||
extern __shared__ int8_t smem[];
|
||||
|
||||
const uint32_t lane_id = get_lane_id();
|
||||
const uint32_t warp_id = get_warp_id();
|
||||
|
||||
// maximize L2 hit rate
|
||||
const uint32_t batch_id = blockIdx.z;
|
||||
const uint32_t bx = blockIdx.x;
|
||||
const uint32_t num_qo_heads = gridDim.y;
|
||||
const uint32_t head_id = blockIdx.y;
|
||||
|
||||
// transfer to base 2 instead of base e with better numerical efficiency
|
||||
sm_scale *= math::log2e;
|
||||
|
||||
// RS holds the fragment of S
|
||||
int32_t RS[num_tiles_q][num_tiles_k][8];
|
||||
DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
|
||||
float m[num_tiles_q][2]; // max
|
||||
float d[num_tiles_q][2]; // denominator
|
||||
|
||||
uint32_t q_scale_idx, k_scale_idx;
|
||||
|
||||
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
|
||||
const uint32_t num_block_q = gridDim.x;
|
||||
q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
|
||||
} else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
|
||||
const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
|
||||
q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
|
||||
} else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
|
||||
const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
|
||||
q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
|
||||
}
|
||||
|
||||
if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
|
||||
const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
|
||||
} else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
|
||||
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
|
||||
} else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
|
||||
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
|
||||
}
|
||||
|
||||
constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
|
||||
: (CTA_K / WARP_K) * 4;
|
||||
|
||||
// initialize o, m, d
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RO[fq][fv][k] = 0.0f;
|
||||
}
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 4; k++) {
|
||||
((int32_t*)RO[fq][fv])[k] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; k++) {
|
||||
m[fq][k] = -5000000.0f;
|
||||
d[fq][k] = 1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr uint32_t K_smem_idx_offset = CTA_Q;
|
||||
constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
|
||||
|
||||
constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
|
||||
: SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
|
||||
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
|
||||
constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
|
||||
constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
|
||||
|
||||
constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
|
||||
: 8;
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
|
||||
: 4;
|
||||
constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8;
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4;
|
||||
constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8;
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4;
|
||||
|
||||
constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
|
||||
constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
|
||||
constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
|
||||
constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
|
||||
constexpr uint32_t V_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V);
|
||||
constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
|
||||
constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
|
||||
|
||||
int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
|
||||
int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
|
||||
half* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
|
||||
uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
|
||||
uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
|
||||
uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
|
||||
|
||||
uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
|
||||
uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
|
||||
uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
|
||||
|
||||
// for causal masking
|
||||
uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
|
||||
uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
|
||||
|
||||
// for loading
|
||||
uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
|
||||
uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
|
||||
uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V;
|
||||
|
||||
const uint32_t num_iterations = SAGE_DIV_CEIL(
|
||||
mask_mode == MaskMode::kCausal
|
||||
? min(kv_len, (bx + 1) * CTA_Q)
|
||||
: kv_len,
|
||||
CTA_K);
|
||||
|
||||
// load Q with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
|
||||
&Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
|
||||
cp_async::commit_group();
|
||||
cp_async::wait_group<0>();
|
||||
__syncthreads();
|
||||
|
||||
// for num_tiles_qk_inner = 1, we load all Qs in register
|
||||
uint32_t RQ[num_tiles_q][4];
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
|
||||
Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
|
||||
}
|
||||
}
|
||||
|
||||
// load K with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
|
||||
float q_scale = Q_scale[q_scale_idx];
|
||||
|
||||
float original_sm_scale = sm_scale;
|
||||
float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
|
||||
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// load V with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
V_load_idx_lane_base += CTA_K;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// do not apply causal mask and out of bound mask for these iterations
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
|
||||
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// load K
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
|
||||
cp_async::commit_group();
|
||||
|
||||
dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
} else {
|
||||
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
// load V
|
||||
load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V);
|
||||
cp_async::commit_group();
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
V_load_idx_lane_base += CTA_K;
|
||||
}
|
||||
|
||||
// second last iter, apply causal mask
|
||||
if (num_iterations > 1) {
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (mask_mode == MaskMode::kCausal) {
|
||||
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
|
||||
}
|
||||
// apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
|
||||
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// load K with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
|
||||
dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
} else {
|
||||
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
// load V with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
V_load_idx_lane_base += CTA_K;
|
||||
}
|
||||
|
||||
// last iter, apply causal mask and out of bound mask
|
||||
{
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (mask_mode == MaskMode::kCausal) {
|
||||
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
|
||||
}
|
||||
// check out of bound in the last iter
|
||||
apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
|
||||
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
|
||||
}
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<0>();
|
||||
__syncthreads();
|
||||
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
} else {
|
||||
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// TODO: thread block sync mdo state for num_warps_k > 0
|
||||
|
||||
normalize_d<num_tiles_q, num_tiles_v, DenominatorAccumUnit>(RO, m, d);
|
||||
|
||||
// save the result
|
||||
// if (get_warp_idx_k<num_warps_q, num_warps_k>() == 0)
|
||||
// {
|
||||
|
||||
// convert half to bfloat16
|
||||
if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::value) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0]));
|
||||
((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1]));
|
||||
((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2]));
|
||||
((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// add v_mean
|
||||
if constexpr (fuse_v_mean) {
|
||||
DTypeOut2 v_mean[2];
|
||||
DTypeOut* V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2;
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16));
|
||||
v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16));
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]);
|
||||
((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]);
|
||||
((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]);
|
||||
((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// save the result to shared memory
|
||||
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
// convert RO to half
|
||||
uint32_t RO_f16[4];
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 4; k++) {
|
||||
if constexpr (std::is_same<DTypeOut, half>::value) {
|
||||
((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
|
||||
} else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::value) {
|
||||
((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
|
||||
}
|
||||
}
|
||||
|
||||
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0];
|
||||
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
|
||||
|
||||
// ! permuted, make sure you know what you are doing
|
||||
((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2];
|
||||
((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0];
|
||||
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1];
|
||||
|
||||
// ! permuted, make sure you know what you are doing
|
||||
((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[2];
|
||||
((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ! do we need to sync here?
|
||||
__syncwarp();
|
||||
|
||||
// shared memory to global memory
|
||||
DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
|
||||
uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
|
||||
uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < O_smem_iters_col; i++) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < O_smem_iters_row; j++) {
|
||||
if (O_load_idx_lane_base < qo_len) {
|
||||
smem_O.store_128b(offset_O, O_lane_ptr);
|
||||
}
|
||||
O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
|
||||
offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
|
||||
}
|
||||
|
||||
offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
|
||||
O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
|
||||
O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
|
||||
}
|
||||
|
||||
if constexpr (return_lse) {
|
||||
uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
|
||||
float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
|
||||
uint32_t fq = (lane_id % 4) / 2;
|
||||
uint32_t k = (lane_id % 4) % 2;
|
||||
|
||||
if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) {
|
||||
lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
|
||||
}
|
||||
}
|
||||
|
||||
// }
|
||||
#else
|
||||
RUNTIME_ASSERT("SageAttention requires SM80 or newer");
|
||||
#endif
|
||||
}
|
||||
|
||||
#undef PACK_SIZE_QK
|
||||
#undef PACK_SIZE_V
|
||||
#undef PACK_SIZE_O
|
||||
#undef MMA_QK_M
|
||||
#undef MMA_QK_N
|
||||
#undef MMA_QK_K
|
||||
#undef MMA_SV_M
|
||||
#undef MMA_SV_N
|
||||
#undef MMA_SV_K
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,622 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
|
||||
|
||||
#include "attn_utils.cuh"
|
||||
#include "math.cuh"
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
#define PACK_SIZE_QK 16 // as if it is int8
|
||||
#define PACK_SIZE_V 16 // fp8
|
||||
#define PACK_SIZE_O 8 // fp16
|
||||
|
||||
// treat as if int8 tensor core
|
||||
#define MMA_QK_M 16
|
||||
#define MMA_QK_N 16
|
||||
#define MMA_QK_K 32
|
||||
|
||||
// fp8 tensor core
|
||||
#define MMA_SV_M 16
|
||||
#define MMA_SV_N 16
|
||||
#define MMA_SV_K 32
|
||||
|
||||
template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale = false, bool fuse_v_mean = false, bool use_pv_fp16_accu = false>
|
||||
__global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
|
||||
// Keep unsupported device passes from parsing architecture-specific intrinsics.
|
||||
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890
|
||||
// compile time check
|
||||
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
|
||||
static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
|
||||
static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
|
||||
static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
|
||||
static_assert(std::is_same<DTypeSVAccum, float>::value, "DTypeSVAccum must be float, half is WIP");
|
||||
static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
|
||||
static_assert(CTA_K % 64 == 0);
|
||||
static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation
|
||||
|
||||
constexpr uint32_t num_warps_q = CTA_Q / WARP_Q;
|
||||
constexpr uint32_t num_warps_k = CTA_K / WARP_K;
|
||||
constexpr uint32_t num_warps = num_warps_q * num_warps_k;
|
||||
constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M;
|
||||
constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N;
|
||||
constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
|
||||
constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N;
|
||||
|
||||
constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
|
||||
constexpr uint32_t O_SMEM_STRIDE = head_dim;
|
||||
// for fp16: head_dim
|
||||
constexpr uint32_t V_SMEM_STRIDE = CTA_K;
|
||||
|
||||
extern __shared__ int8_t smem[];
|
||||
|
||||
const uint32_t lane_id = get_lane_id();
|
||||
const uint32_t warp_id = get_warp_id();
|
||||
|
||||
// maximize L2 hit rate
|
||||
const uint32_t batch_id = blockIdx.z;
|
||||
const uint32_t bx = blockIdx.x;
|
||||
const uint32_t num_qo_heads = gridDim.y;
|
||||
const uint32_t head_id = blockIdx.y;
|
||||
|
||||
// transfer to base 2 instead of base e with better numerical efficiency
|
||||
sm_scale *= math::log2e;
|
||||
|
||||
// RS holds the fragment of S
|
||||
int32_t RS[num_tiles_q][num_tiles_k][8];
|
||||
DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
|
||||
float m[num_tiles_q][2]; // max
|
||||
float d[num_tiles_q][2]; // denominator
|
||||
|
||||
uint32_t q_scale_idx, k_scale_idx;
|
||||
|
||||
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
|
||||
const uint32_t num_block_q = gridDim.x;
|
||||
q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
|
||||
} else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
|
||||
const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
|
||||
q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
|
||||
} else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
|
||||
const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
|
||||
q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
|
||||
}
|
||||
|
||||
if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
|
||||
const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
|
||||
} else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
|
||||
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
|
||||
} else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
|
||||
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
|
||||
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
|
||||
}
|
||||
|
||||
constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
|
||||
: (CTA_K / WARP_K) * 4;
|
||||
|
||||
// initialize o, m, d
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RO[fq][fv][k] = 0.0f;
|
||||
}
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 4; k++) {
|
||||
((int32_t*)RO[fq][fv])[k] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; k++) {
|
||||
m[fq][k] = -5000000.0f;
|
||||
d[fq][k] = 1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr uint32_t K_smem_idx_offset = CTA_Q;
|
||||
constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
|
||||
|
||||
constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
|
||||
: SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
|
||||
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
|
||||
// for fp16: 32
|
||||
constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
|
||||
constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
|
||||
smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
|
||||
|
||||
constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
|
||||
: 8;
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
|
||||
: 4;
|
||||
// for fp16: 32
|
||||
constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8;
|
||||
// for fp16: 32
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4;
|
||||
constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8;
|
||||
constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4;
|
||||
|
||||
constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
|
||||
constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
|
||||
constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
|
||||
constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
|
||||
// for fp16: CTA_K
|
||||
constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V);
|
||||
constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
|
||||
constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
|
||||
|
||||
int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
|
||||
int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
|
||||
// for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v
|
||||
int8_t* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
|
||||
uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
|
||||
uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
|
||||
uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
|
||||
|
||||
uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
|
||||
uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
|
||||
// for fp 16:
|
||||
// uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
|
||||
uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2);
|
||||
|
||||
// for causal masking
|
||||
uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
|
||||
uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
|
||||
|
||||
// for loading
|
||||
uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
|
||||
uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
|
||||
|
||||
const uint32_t num_iterations = SAGE_DIV_CEIL(
|
||||
mask_mode == MaskMode::kCausal
|
||||
? min(kv_len, (bx + 1) * CTA_Q)
|
||||
: kv_len,
|
||||
CTA_K);
|
||||
|
||||
// load Q with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
|
||||
&Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
|
||||
cp_async::commit_group();
|
||||
cp_async::wait_group<0>();
|
||||
__syncthreads();
|
||||
|
||||
// for num_tiles_qk_inner = 1, we load all Qs in register
|
||||
uint32_t RQ[num_tiles_q][4];
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
|
||||
Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
|
||||
}
|
||||
}
|
||||
|
||||
// load K with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
|
||||
float q_scale = Q_scale[q_scale_idx];
|
||||
|
||||
float original_sm_scale = sm_scale;
|
||||
float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
|
||||
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// load V
|
||||
// ! we assume that V is padded. If not, there might be illegal memory access or nan issue.
|
||||
// for fp16:
|
||||
// load_global_to_share stride_seq_v
|
||||
load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
|
||||
cp_async::commit_group();
|
||||
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
|
||||
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// load K without predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
|
||||
cp_async::commit_group();
|
||||
|
||||
dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// for fp16:
|
||||
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
// smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
if constexpr (!use_pv_fp16_accu) {
|
||||
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
// load V
|
||||
// for fp16:
|
||||
// load_global_to_share stride_seq_v
|
||||
load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
|
||||
cp_async::commit_group();
|
||||
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
}
|
||||
|
||||
// second last iter, apply causal mask
|
||||
if (num_iterations > 1) {
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (mask_mode == MaskMode::kCausal) {
|
||||
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
|
||||
}
|
||||
// apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
|
||||
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// load K with predicate
|
||||
load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
|
||||
&K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
|
||||
cp_async::commit_group();
|
||||
|
||||
dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
|
||||
sm_scale = original_sm_scale * dequant_scale;
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// for fp16:
|
||||
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
// smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
if constexpr (!use_pv_fp16_accu) {
|
||||
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
// load V
|
||||
// for fp16:
|
||||
// load_global_to_share stride_seq_v
|
||||
load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
|
||||
&V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
|
||||
cp_async::commit_group();
|
||||
K_load_idx_lane_base += CTA_K;
|
||||
}
|
||||
|
||||
// last iter, apply causal mask and out of bound mask
|
||||
{
|
||||
// ensure K is ready
|
||||
cp_async::wait_group<1>();
|
||||
__syncthreads();
|
||||
|
||||
// compute QK^T
|
||||
if constexpr (num_tiles_qk_inner == 1) {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_K, RS, RQ, K_smem_offset_mma);
|
||||
} else {
|
||||
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
|
||||
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
|
||||
}
|
||||
|
||||
float RS_f32[num_tiles_q][num_tiles_k][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (mask_mode == MaskMode::kCausal) {
|
||||
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
|
||||
}
|
||||
apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
|
||||
K_idx_lane_base += CTA_K;
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
|
||||
}
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
|
||||
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
|
||||
}
|
||||
|
||||
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
|
||||
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
|
||||
|
||||
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
|
||||
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
|
||||
}
|
||||
|
||||
// ensure V is ready
|
||||
cp_async::wait_group<0>();
|
||||
__syncthreads();
|
||||
|
||||
// for fp16:
|
||||
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
|
||||
// smem_V, RS_f16, RO, d, V_smem_offset_mma);
|
||||
if constexpr (!use_inst_buffer) {
|
||||
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
if constexpr (!use_pv_fp16_accu) {
|
||||
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
} else {
|
||||
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
|
||||
smem_V, RS_f8, RO, d);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving.
|
||||
|
||||
normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d);
|
||||
|
||||
// ! here we just implement the case for fp32 acumulation
|
||||
if constexpr (fuse_v_scale) {
|
||||
float v_scale[4];
|
||||
float* V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16));
|
||||
((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8));
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
RO[fq][fv][0] *= v_scale[0];
|
||||
RO[fq][fv][1] *= v_scale[1];
|
||||
RO[fq][fv][2] *= v_scale[0];
|
||||
RO[fq][fv][3] *= v_scale[1];
|
||||
RO[fq][fv][4] *= v_scale[2];
|
||||
RO[fq][fv][5] *= v_scale[3];
|
||||
RO[fq][fv][6] *= v_scale[2];
|
||||
RO[fq][fv][7] *= v_scale[3];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (fuse_v_mean) {
|
||||
float v_mean[4];
|
||||
float* V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16));
|
||||
((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8));
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
RO[fq][fv][0] += v_mean[0];
|
||||
RO[fq][fv][1] += v_mean[1];
|
||||
RO[fq][fv][2] += v_mean[0];
|
||||
RO[fq][fv][3] += v_mean[1];
|
||||
RO[fq][fv][4] += v_mean[2];
|
||||
RO[fq][fv][5] += v_mean[3];
|
||||
RO[fq][fv][6] += v_mean[2];
|
||||
RO[fq][fv][7] += v_mean[3];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// save the result to shared memory
|
||||
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
|
||||
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
// convert RO to half
|
||||
uint32_t RO_f16[4];
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 4; k++) {
|
||||
if constexpr (std::is_same<DTypeOut, half>::value) {
|
||||
((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
|
||||
} else {
|
||||
((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
|
||||
}
|
||||
}
|
||||
|
||||
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0];
|
||||
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
|
||||
|
||||
offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1);
|
||||
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2];
|
||||
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// TODO: not implement
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ! do we need to sync here?
|
||||
__syncwarp();
|
||||
|
||||
// shared memory to global memory
|
||||
DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
|
||||
uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
|
||||
uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < O_smem_iters_col; i++) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < O_smem_iters_row; j++) {
|
||||
if (O_load_idx_lane_base < qo_len) {
|
||||
smem_O.store_128b(offset_O, O_lane_ptr);
|
||||
}
|
||||
O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
|
||||
offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
|
||||
}
|
||||
|
||||
offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
|
||||
O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
|
||||
O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
|
||||
}
|
||||
|
||||
if constexpr (return_lse) {
|
||||
// ! this only works for num_tiles_q = 2
|
||||
uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
|
||||
float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
|
||||
uint32_t fq = (lane_id % 4) / 2;
|
||||
uint32_t k = (lane_id % 4) % 2;
|
||||
|
||||
if (lse_idx < qo_len) {
|
||||
lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
|
||||
}
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer");
|
||||
#endif
|
||||
}
|
||||
|
||||
#undef PACK_SIZE_QK
|
||||
#undef PACK_SIZE_V
|
||||
#undef PACK_SIZE_O
|
||||
#undef MMA_QK_M
|
||||
#undef MMA_QK_N
|
||||
#undef MMA_QK_K
|
||||
#undef MMA_SV_M
|
||||
#undef MMA_SV_N
|
||||
#undef MMA_SV_K
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,814 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_pipeline_primitives.h>
|
||||
|
||||
#include "cp_async.cuh"
|
||||
#include "math.cuh"
|
||||
#include "mma.cuh"
|
||||
#include "numeric_conversion.cuh"
|
||||
#include "permuted_smem.cuh"
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
#define WARP_SIZE 32
|
||||
|
||||
#define S_FP8_OFFSET 8.807f
|
||||
#define S_FP8_OFFSET_EXP 6680.8477f
|
||||
#define S_FP8_OFFSET_EXP_INV 0.0022326917f
|
||||
|
||||
#define SAGE_DIV_CEIL(M, N) (((M) + (N) - 1) / (N))
|
||||
|
||||
enum class MaskMode {
|
||||
kNone = 0,
|
||||
kCausal = 1,
|
||||
};
|
||||
|
||||
enum class DataType {
|
||||
kHalf,
|
||||
kInt8,
|
||||
kInt4,
|
||||
kE4M3,
|
||||
kE5M2,
|
||||
};
|
||||
|
||||
enum class QuantGranularity {
|
||||
kPerTensor = 0,
|
||||
kPerBlock = 1,
|
||||
kPerWarp = 2,
|
||||
kPerThread = 3,
|
||||
};
|
||||
|
||||
enum class ComputeUnit {
|
||||
kTensorCore,
|
||||
kCudaCore,
|
||||
};
|
||||
|
||||
__device__ __forceinline__ uint32_t get_warp_id() {
|
||||
return threadIdx.y;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t get_lane_id() {
|
||||
return threadIdx.x;
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k>
|
||||
__device__ __forceinline__ uint32_t get_warp_idx_q() {
|
||||
return get_warp_id() / num_warps_k;
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k>
|
||||
__device__ __forceinline__ uint32_t get_warp_idx_k() {
|
||||
return get_warp_id() % num_warps_k;
|
||||
}
|
||||
|
||||
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
|
||||
__device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
|
||||
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
|
||||
static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
|
||||
|
||||
constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < smem_iters_col; i++) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < smem_iters_row; j++) {
|
||||
smem.load_128b_async(smem_offset, *lane_ptr);
|
||||
*lane_ptr += (global_to_shared_line_lanes * pack_size);
|
||||
smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
|
||||
}
|
||||
|
||||
smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
|
||||
*lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
|
||||
}
|
||||
smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
|
||||
*lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
|
||||
}
|
||||
|
||||
// with predicate
|
||||
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
|
||||
__device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem, uint32_t base_idx, uint32_t max_len) {
|
||||
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
|
||||
static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
|
||||
|
||||
constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < smem_iters_col; i++) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < smem_iters_row; j++) {
|
||||
// Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN.
|
||||
smem.load_128b_async<std::is_same<T, half>::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len);
|
||||
*lane_ptr += (global_to_shared_line_lanes * pack_size);
|
||||
smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
|
||||
}
|
||||
|
||||
smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
|
||||
*lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
|
||||
base_idx += global_to_shared_copy_lines_per_warp_per_iter;
|
||||
}
|
||||
smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
|
||||
*lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
|
||||
}
|
||||
|
||||
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA>
|
||||
__device__ __forceinline__ void load_fp8_V_global_to_share(int8_t** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
|
||||
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
|
||||
constexpr uint32_t pack_size_fp8 = 16;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < smem_iters_col; i++) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < smem_iters_row; j++) {
|
||||
smem.load_128b_async(smem_offset, *lane_ptr);
|
||||
*lane_ptr += (global_to_shared_line_lanes * pack_size_fp8);
|
||||
smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
|
||||
}
|
||||
|
||||
smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
|
||||
*lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8));
|
||||
}
|
||||
smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
|
||||
// for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
|
||||
*lane_ptr += CTA; // ! prevent underflow
|
||||
*lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
|
||||
__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_Q, const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t& offset_Q, uint32_t& offset_K) {
|
||||
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
|
||||
|
||||
uint32_t RQ[num_tiles_q][4];
|
||||
uint32_t RK[4];
|
||||
|
||||
// the first iteration, mma mode is kInit
|
||||
#pragma unroll
|
||||
for (uint32_t iter = 0; iter < 1; iter++) {
|
||||
// load RQ
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
|
||||
offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
|
||||
}
|
||||
// ! using permutation invariance
|
||||
offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
// load RK
|
||||
smem_K.ldmatrix_m8n8x4(offset_K, RK);
|
||||
offset_K = smem_K.advance_offset_by_row<16>(offset_K);
|
||||
|
||||
// mma
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (DTypeQK == DataType::kInt8) {
|
||||
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
|
||||
} else if constexpr (DTypeQK == DataType::kInt4) {
|
||||
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
|
||||
}
|
||||
}
|
||||
}
|
||||
offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
|
||||
}
|
||||
|
||||
// following iteration, mma mode is kInplace
|
||||
#pragma unroll
|
||||
for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) {
|
||||
// load RQ
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
|
||||
offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
|
||||
}
|
||||
offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
// load RK
|
||||
smem_K.ldmatrix_m8n8x4(offset_K, RK);
|
||||
offset_K = smem_K.advance_offset_by_row<16>(offset_K);
|
||||
|
||||
// mma
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (DTypeQK == DataType::kInt8) {
|
||||
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
|
||||
} else if constexpr (DTypeQK == DataType::kInt4) {
|
||||
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
|
||||
}
|
||||
}
|
||||
}
|
||||
offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
|
||||
}
|
||||
|
||||
offset_Q -= (2 * num_tiles_qk_inner);
|
||||
offset_K -= (2 * num_tiles_qk_inner);
|
||||
}
|
||||
|
||||
// for case when num_tiles_qk_inner = 1
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
|
||||
__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) {
|
||||
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
|
||||
static_assert(num_tiles_qk_inner == 1);
|
||||
|
||||
uint32_t RK[4];
|
||||
|
||||
// mma mode is kInit
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
// load RK
|
||||
smem_K.ldmatrix_m8n8x4(offset_K, RK);
|
||||
offset_K = smem_K.advance_offset_by_row<16>(offset_K);
|
||||
|
||||
// mma
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (DTypeQK == DataType::kInt8) {
|
||||
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
|
||||
} else if constexpr (DTypeQK == DataType::kInt4) {
|
||||
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
|
||||
__device__ __forceinline__ void apply_causal_mask(const uint32_t& Q_idx_lane_base, const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2);
|
||||
const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
|
||||
const bool out_of_boundary = (kv_idx > q_idx);
|
||||
|
||||
if constexpr (std::is_same<DTypeQKAccum, float>::value) {
|
||||
RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
|
||||
} else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
|
||||
RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
|
||||
__device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t& kv_len) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
|
||||
const bool out_of_boundary = (kv_idx >= kv_len);
|
||||
|
||||
if constexpr (std::is_same<DTypeQKAccum, float>::value) {
|
||||
RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
|
||||
} else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
|
||||
RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// for DTypeQKAccum float
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale = false, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float& sm_scale) {
|
||||
static_assert(std::is_same<DTypeSVAccum, half>::value || (!use_half_o_scale));
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; k++) {
|
||||
// assign the smallest value possible
|
||||
float m_prev = m[fq][k];
|
||||
float m_temp = -5000000.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]),
|
||||
max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5]));
|
||||
m_temp = max(m_temp, m_local);
|
||||
}
|
||||
|
||||
if constexpr (!fuse_scale) {
|
||||
if constexpr (exp_offset) {
|
||||
m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET);
|
||||
} else {
|
||||
m_temp *= sm_scale;
|
||||
}
|
||||
} else if constexpr (exp_offset) {
|
||||
m_temp += (-S_FP8_OFFSET);
|
||||
}
|
||||
|
||||
// exchange element with the 4 threads in the row
|
||||
m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1)); // 0 exchange with 1, 2 exchange with 3
|
||||
m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2)); // 0 exchange with 2, 1 exchange with 3
|
||||
|
||||
m[fq][k] = max(m[fq][k], m_temp);
|
||||
|
||||
float o_scale = math::ptx_exp2(m_prev - m[fq][k]);
|
||||
|
||||
// update denominator
|
||||
d[fq][k] *= o_scale;
|
||||
|
||||
half2 o_scale2;
|
||||
if constexpr (use_half_o_scale) {
|
||||
o_scale2 = __floats2half2_rn(o_scale, o_scale);
|
||||
}
|
||||
|
||||
// update RO
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
RO[fq][fv][k * 2 + 0] *= o_scale;
|
||||
RO[fq][fv][k * 2 + 1] *= o_scale;
|
||||
RO[fq][fv][k * 2 + 4] *= o_scale;
|
||||
RO[fq][fv][k * 2 + 5] *= o_scale;
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
if constexpr (use_half_o_scale) {
|
||||
((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2);
|
||||
((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2);
|
||||
} else {
|
||||
RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale);
|
||||
RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale);
|
||||
RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale);
|
||||
RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// raise RS to exponent
|
||||
float negative_m = -m[fq][k];
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
if constexpr (fuse_scale) {
|
||||
RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m);
|
||||
RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m);
|
||||
RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m);
|
||||
RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m);
|
||||
} else {
|
||||
RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m));
|
||||
RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m));
|
||||
RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m));
|
||||
RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T>
|
||||
__device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) {
|
||||
static_assert(sizeof(T) == 4);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]);
|
||||
((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]);
|
||||
((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]);
|
||||
((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k>
|
||||
__device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
|
||||
floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4);
|
||||
floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6);
|
||||
floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4);
|
||||
floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k>
|
||||
__device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
|
||||
halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2);
|
||||
halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3);
|
||||
halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2);
|
||||
halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k>
|
||||
__device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
|
||||
e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]);
|
||||
e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1);
|
||||
e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2);
|
||||
e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T>
|
||||
__device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore) ? 4 : 8], float d[][2]) {
|
||||
// for compute unit cuda core, RS is float
|
||||
// for compute unit tensor core, RS is packed half
|
||||
static_assert((std::is_same<T, float>::value && compute_unit == ComputeUnit::kCudaCore) ||
|
||||
(std::is_same<T, uint32_t>::value && compute_unit == ComputeUnit::kTensorCore));
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
if constexpr (compute_unit == ComputeUnit::kTensorCore) {
|
||||
// full accumulate with tensor core
|
||||
mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk]));
|
||||
} else if constexpr (compute_unit == ComputeUnit::kCudaCore) {
|
||||
// partial accumulate with cuda core
|
||||
d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5];
|
||||
d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_k>
|
||||
__device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) {
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
|
||||
mma::rowsum_f8f8f32(d[fq], RS[fq][fk]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
|
||||
uint32_t smem_V_row_base = get_warp_idx_k<num_warps_q, num_warps_k>() * (num_tiles_k * 16) + get_lane_id() % 16;
|
||||
uint32_t smem_V_col_base = get_lane_id() / 16;
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2);
|
||||
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
|
||||
static_assert(sizeof(T) == 4);
|
||||
|
||||
// ! be sure you know what you are doing
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
|
||||
}
|
||||
}
|
||||
|
||||
offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
|
||||
}
|
||||
|
||||
// make offset_V their original value
|
||||
offset_V -= (16 * num_tiles_k * stride);
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
|
||||
static_assert(sizeof(T) == 4);
|
||||
static_assert(std::is_same<DTypeSVAccum, float>::value);
|
||||
|
||||
uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4];
|
||||
|
||||
// ! be sure you know what you are doing
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < 1; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
{
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInit>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
|
||||
}
|
||||
}
|
||||
|
||||
offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 1; fk < num_tiles_k; fk++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
{
|
||||
mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInplaceUpdate>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
|
||||
}
|
||||
}
|
||||
|
||||
offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
|
||||
}
|
||||
|
||||
// accumulate into RO
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x);
|
||||
RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y);
|
||||
RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x);
|
||||
RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y);
|
||||
RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x);
|
||||
RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y);
|
||||
RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x);
|
||||
RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y);
|
||||
}
|
||||
}
|
||||
|
||||
// make offset_V their original value
|
||||
offset_V -= (16 * num_tiles_k * stride);
|
||||
}
|
||||
|
||||
template <uint32_t num_tiles_q, uint32_t num_tiles_v,
|
||||
ComputeUnit compute_unit = ComputeUnit::kTensorCore, // compute unit for accumulate_d
|
||||
typename DTypeQKAccum,
|
||||
typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) {
|
||||
if constexpr (compute_unit == ComputeUnit::kCudaCore) {
|
||||
// accumulate_d performs partial accumulation with cuda core
|
||||
// aggregate d
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; k++) {
|
||||
d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3
|
||||
d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// divide O by d
|
||||
float d_rcp[num_tiles_q][2];
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 2; k++) {
|
||||
// TODO: check m to prevent nan
|
||||
d_rcp[fq][k] = math::ptx_rcp(d[fq][k]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < 8; k++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2];
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
|
||||
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
|
||||
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
|
||||
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
|
||||
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
|
||||
smem_V.ldmatrix_m8n8x4(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// ! Not Implemented
|
||||
}
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
|
||||
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
|
||||
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
|
||||
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
|
||||
|
||||
float RO_inst_buf[num_tiles_q][num_tiles_v][8];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < 1; fk++) {
|
||||
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
|
||||
smem_V.ldmatrix_m8n8x4(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// ! Not Implemented
|
||||
}
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
|
||||
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
|
||||
smem_V.ldmatrix_m8n8x4(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// ! Not Implemented
|
||||
}
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
|
||||
RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
|
||||
RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
|
||||
RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
|
||||
RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
|
||||
RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
|
||||
RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
|
||||
RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
|
||||
__device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
|
||||
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
|
||||
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
|
||||
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
|
||||
|
||||
uint32_t RO_int32[num_tiles_q][num_tiles_v][4];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 0; fk < 1; fk++) {
|
||||
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
|
||||
smem_V.ldmatrix_m8n8x4(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
// mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
|
||||
mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// ! Not Implemented
|
||||
}
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
|
||||
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
|
||||
#pragma unroll
|
||||
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
|
||||
// load RV
|
||||
uint32_t RV[4];
|
||||
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
|
||||
smem_V.ldmatrix_m8n8x4(offset_V, RV);
|
||||
#pragma unroll
|
||||
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
|
||||
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
|
||||
// mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
|
||||
mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
|
||||
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
|
||||
// ! Not Implemented
|
||||
}
|
||||
}
|
||||
offset_V = smem_V.advance_offset_by_row<16>(offset_V);
|
||||
}
|
||||
}
|
||||
float RO_tmp_float[2];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < num_tiles_q; i++) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < num_tiles_v; j++) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; k++) {
|
||||
unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]);
|
||||
RO[i][j][k * 2 + 0] += RO_tmp_float[0];
|
||||
RO[i][j][k * 2 + 1] += RO_tmp_float[1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// #pragma unroll
|
||||
// for (uint32_t fq = 0; fq < num_tiles_q; fq++)
|
||||
// {
|
||||
// #pragma unroll
|
||||
// for (uint32_t fv = 0; fv < num_tiles_v; fv++)
|
||||
// {
|
||||
// RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
|
||||
// RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
|
||||
// RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
|
||||
// RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
|
||||
// RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
|
||||
// RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
|
||||
// RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
|
||||
// RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,148 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/cp_async.cuh
|
||||
* Copyright (c) 2023 by FlashInfer team.
|
||||
* Small modifications made by SageAttention team, 2024 (e.g., renamed namespace).
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
namespace cp_async {
|
||||
|
||||
enum class SharedMemFillMode {
|
||||
kFillZero, // Fill zero to shared memory when predicate is false
|
||||
kNoFill // Do not fill zero to shared memory when predicate is false
|
||||
};
|
||||
|
||||
enum class PrefetchMode {
|
||||
kNoPrefetch, // Do not fetch additional data from global memory to L2
|
||||
kPrefetch // Fetch additional data from global memory to L2
|
||||
};
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ >= 11)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
|
||||
#define CP_ASYNC_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
|
||||
* cp.async instructions to a group
|
||||
*/
|
||||
__device__ __forceinline__ void commit_group() {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.wait_group instruction
|
||||
* \tparam n Wait till most recent n groups are committed
|
||||
*/
|
||||
template <size_t n>
|
||||
__device__ __forceinline__ void wait_group() {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory
|
||||
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
|
||||
* \tparam T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, typename T>
|
||||
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
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));
|
||||
} else {
|
||||
asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(gmem_ptr), "n"(16), "r"(16));
|
||||
}
|
||||
#else
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
|
||||
* global memory to shared memory with predicate.
|
||||
* \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 T Data type
|
||||
* \param smem_ptr Pointer to shared memory
|
||||
* \param gmem_ptr Pointer to global memory
|
||||
* \param predicate Predicate value
|
||||
* \note fill zero is slower than not fill zero
|
||||
*/
|
||||
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
|
||||
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
|
||||
#ifdef CP_ASYNC_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
int src_in_bytes = predicate ? 16 : 0;
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
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));
|
||||
} else {
|
||||
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));
|
||||
}
|
||||
} else {
|
||||
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
} else {
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)predicate),
|
||||
"r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
|
||||
}
|
||||
}
|
||||
#else
|
||||
if (predicate) {
|
||||
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
|
||||
} else {
|
||||
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
|
||||
*((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace cp_async
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,166 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/math.cuh
|
||||
* Copyright (c) 2023 by FlashInfer team.
|
||||
* Small modifications made by SageAttention team, 2024 (e.g., renamed namespace).
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_MATH_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_MATH_CUH
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace sageattention {
|
||||
#ifndef USHORT_TYPE
|
||||
#define USHORT_TYPE
|
||||
typedef unsigned short ushort;
|
||||
#endif
|
||||
|
||||
namespace math {
|
||||
|
||||
// log2(e)
|
||||
constexpr float log2e = 1.44269504088896340736f;
|
||||
constexpr float log2e_recp = 1.0f / log2e;
|
||||
|
||||
__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;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ float ptx_exp2(float x) {
|
||||
float y;
|
||||
asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ float ptx_log2(float x) {
|
||||
float y;
|
||||
asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ half2 ptx_exp2(half2 x) {
|
||||
uint32_t y_u32;
|
||||
uint32_t x_u32 = half2_as_uint32(x);
|
||||
asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
|
||||
return uint32_as_half2(y_u32);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ half ptx_exp2(half x) {
|
||||
ushort y_u16;
|
||||
asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
|
||||
return __ushort_as_half(y_u16);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ float ptx_rcp(float x) {
|
||||
float y;
|
||||
asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
|
||||
* between threads in a warp.
|
||||
* \param x The value in the source lane
|
||||
* \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) {
|
||||
float y;
|
||||
asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
|
||||
: "=f"(y)
|
||||
: "f"(x), "r"(lane_mask));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
|
||||
* shuffle between threads in a warp.
|
||||
* \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]
|
||||
*/
|
||||
__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
|
||||
return __shfl_xor_sync(0xffffffff, x, lane_mask);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ float rsqrt(float x) {
|
||||
float y;
|
||||
asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ float tanh(float x) {
|
||||
float y;
|
||||
asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
|
||||
return y;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ half2 tanh(half2 x) {
|
||||
uint32_t y_u32;
|
||||
uint32_t x_u32 = half2_as_uint32(x);
|
||||
asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
|
||||
return uint32_as_half2(y_u32);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
|
||||
* \param x input
|
||||
*/
|
||||
__forceinline__ __device__ half tanh(half x) {
|
||||
ushort y_u16;
|
||||
asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
|
||||
return __ushort_as_half(y_u16);
|
||||
}
|
||||
|
||||
} // namespace math
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,683 @@
|
||||
/*
|
||||
* Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/mma.cuh
|
||||
* Copyright (c) 2023 by FlashInfer team.
|
||||
*
|
||||
* Modifications copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_MMA_CUH
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
namespace mma {
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ >= 11)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
|
||||
#define MMA_F16F16F32_M16N8K16_ENABLED
|
||||
#define MMA_F16F16F16_M16N8K16_ENABLED
|
||||
#define MMA_S8S8S32_M16N8K32_ENABLED
|
||||
#define MMA_S4S4S32_M16N8K64_ENABLED
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750))
|
||||
#define MMA_F16F16F32_M16N8K8_ENABLED
|
||||
#define MMA_F16F16F16_M16N8K8_ENABLED
|
||||
#define LDMATRIX_M8N8X2_ENABLED
|
||||
#define LDMATRIX_M8N8X4_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
|
||||
#define MMA_F8F8F32_M16N8K16_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
|
||||
#define MMA_F8F8F16_M16N8K16_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__)
|
||||
#define RUNTIME_ASSERT(x) __brkpt()
|
||||
#else
|
||||
|
||||
#define RUNTIME_ASSERT(x) assert(0 && x)
|
||||
#endif
|
||||
|
||||
enum class MMAMode {
|
||||
kInit = 0U,
|
||||
kInplaceUpdate = 1U,
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory
|
||||
* to fragment
|
||||
* \tparam T data type of the fragment
|
||||
* \param R pointer to the fragment
|
||||
* \param smem_ptr pointer to the shared memory
|
||||
*/
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
|
||||
#ifdef LDMATRIX_M8N8X2_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
|
||||
: "=r"(R[0]), "=r"(R[1])
|
||||
: "r"(smem_int_ptr));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory
|
||||
* to fragment
|
||||
* \tparam T data type of the fragment
|
||||
* \param R pointer to the fragment
|
||||
* \param smem_ptr pointer to the shared memory
|
||||
*/
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
|
||||
#ifdef LDMATRIX_M8N8X4_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
|
||||
: "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
|
||||
: "r"(smem_int_ptr));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from
|
||||
* shared memory to fragment and transposes the fragment
|
||||
* \tparam T data type of the fragment
|
||||
* \param R pointer to the fragment
|
||||
* \param smem_ptr pointer to the shared memory
|
||||
*/
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
|
||||
#ifdef LDMATRIX_M8N8X4_ENABLED
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
|
||||
: "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
|
||||
: "r"(smem_int_ptr));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
|
||||
* multiplication, accumulated in f32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
|
||||
// ! only support half dtype now
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
|
||||
"f"(C[2]), "f"(C[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
|
||||
* multiplication, accumulated in f32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
|
||||
// ! only support half dtype now
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
|
||||
"f"(C[2]), "f"(C[3]));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
|
||||
"f"(C[6]), "f"(C[7]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
|
||||
* multiplication, accumulated in f16.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[0]), "=r"(C[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[0]), "=r"(C[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
|
||||
* multiplication, accumulated in f16.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[0]), "=r"(C[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[0]), "=r"(C[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix
|
||||
* multiplication, accumulated in int32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
|
||||
"r"(C[2]), "r"(C[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix
|
||||
* multiplication, accumulated in int32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
|
||||
"r"(C[2]), "r"(C[3]));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
|
||||
"r"(C[6]), "r"(C[7]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix
|
||||
* multiplication, accumulated in int32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
|
||||
"r"(C[2]), "r"(C[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix
|
||||
* multiplication, accumulated in int32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
|
||||
"r"(C[2]), "r"(C[3]));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
|
||||
"r"(C[6]), "r"(C[7]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
|
||||
"r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix
|
||||
* multiplication, accumulated in fp32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
|
||||
"f"(C[2]), "f"(C[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
|
||||
* multiplication, accumulated in fp16.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
|
||||
// uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
|
||||
#ifdef MMA_F8F8F16_M16N8K16_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C_uint32[0]), "=r"(C_uint32[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1]));
|
||||
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C_uint32[2]), "=r"(C_uint32[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C_uint32[0]), "=r"(C_uint32[1])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
|
||||
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
|
||||
"{%0, %1},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, %9};\n"
|
||||
: "=r"(C_uint32[2]), "=r"(C_uint32[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
|
||||
* multiplication, accumulated in fp32.
|
||||
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
|
||||
* \param C pointer to the accumulator
|
||||
* \param A pointer to the fragment of matrix A
|
||||
* \param B pointer to the fragment of matrix B
|
||||
*/
|
||||
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
|
||||
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
|
||||
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
|
||||
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
|
||||
"f"(C[2]), "f"(C[3]));
|
||||
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
|
||||
"f"(C[6]), "f"(C[7]));
|
||||
} else if constexpr (mma_mode == MMAMode::kInit) {
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
|
||||
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
|
||||
"f"(0.f), "f"(0.f));
|
||||
}
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Use mma instructions to compute rowsum.
|
||||
*/
|
||||
__device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
|
||||
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
|
||||
"{%0, _, %1, _},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, 0., %9, 0.};\n"
|
||||
"}\n"
|
||||
: "=f"(d[0]), "=f"(d[1])
|
||||
: "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320), // 1006648320 packs two 1.0f in half precision
|
||||
"r"(1006648320), "f"(d[0]), "f"(d[1]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Use mma instructions to compute rowsum.
|
||||
*/
|
||||
__device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
|
||||
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0, _, %1, _},"
|
||||
"{%2, %3, %4, %5},"
|
||||
"{%6, %7},"
|
||||
"{%8, 0., %9, 0.};\n"
|
||||
: "=f"(d[0]), "=f"(d[1])
|
||||
: "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504), // 943208504 packs four 1.0f in e4m3
|
||||
"f"(d[0]), "f"(d[1]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace mma
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,148 @@
|
||||
/*
|
||||
* Copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Inspired by CUTLASS, https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/numeric_conversion.h
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
|
||||
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
|
||||
#define FP8_CAST_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__)
|
||||
#define RUNTIME_ASSERT(x) __brkpt()
|
||||
#else
|
||||
|
||||
#define RUNTIME_ASSERT(x) assert(0 && x)
|
||||
#endif
|
||||
|
||||
__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
|
||||
uint16_t h0 = source & 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[1]) : "h"(h1));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo;\n"
|
||||
".reg .b16 hi;\n"
|
||||
"cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
|
||||
"cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
|
||||
"mov.b32 %0, {lo, hi};\n"
|
||||
"}"
|
||||
: "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo;\n"
|
||||
".reg .b16 hi;\n"
|
||||
"cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n"
|
||||
"cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n"
|
||||
"mov.b32 %0, {lo, hi};\n"
|
||||
"}"
|
||||
: "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo;\n"
|
||||
".reg .b16 hi;\n"
|
||||
"cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n"
|
||||
"cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n"
|
||||
"mov.b32 %0, {lo, hi};\n"
|
||||
"}"
|
||||
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo;\n"
|
||||
".reg .b16 hi;\n"
|
||||
"cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n"
|
||||
"cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n"
|
||||
"mov.b32 %0, {lo, hi};\n"
|
||||
"}"
|
||||
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo, hi;\n"
|
||||
"mov.b32 {lo, hi}, %2;\n"
|
||||
"cvt.rn.f16x2.e4m3x2 %0, lo;\n"
|
||||
"cvt.rn.f16x2.e4m3x2 %1, hi;\n"
|
||||
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
|
||||
#ifdef FP8_CAST_ENABLED
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b16 lo, hi;\n"
|
||||
"mov.b32 {lo, hi}, %2;\n"
|
||||
"cvt.rn.f16x2.e5m2x2 %0, lo;\n"
|
||||
"cvt.rn.f16x2.e5m2x2 %1, hi;\n"
|
||||
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
|
||||
#else
|
||||
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int8_t float_to_int8_rn(float x) {
|
||||
uint32_t dst;
|
||||
asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
|
||||
return reinterpret_cast<const int8_t&>(dst);
|
||||
}
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,201 @@
|
||||
/*
|
||||
* Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/permuted_smem.cuh
|
||||
* Copyright (c) 2023 by FlashInfer team.
|
||||
*
|
||||
* Modifications copyright (c) 2024 by SageAttention team.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH
|
||||
#define SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "cp_async.cuh"
|
||||
#include "mma.cuh"
|
||||
|
||||
namespace sageattention {
|
||||
|
||||
enum class SwizzleMode {
|
||||
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)
|
||||
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
|
||||
using b128_t = uint4;
|
||||
|
||||
/*!
|
||||
* \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
|
||||
* 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,
|
||||
* 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 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.
|
||||
* Under this assumption, we do not permute for k32B mode.
|
||||
*/
|
||||
template <SwizzleMode swizzle_mode, uint32_t stride>
|
||||
struct smem_t {
|
||||
// The base pointer.
|
||||
b128_t* base;
|
||||
// How many b128_t value a row contains
|
||||
// uint32_t stride;
|
||||
|
||||
__device__ __forceinline__ smem_t() : base(nullptr) {}
|
||||
template <typename T>
|
||||
__device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
|
||||
if constexpr (swizzle_mode == SwizzleMode::k128B) {
|
||||
static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
|
||||
static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k32B) {
|
||||
static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
|
||||
} else {
|
||||
static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Set the base pointer.
|
||||
*/
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void set_base(T* new_base) {
|
||||
base = (b128_t*)new_base;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Compute the element offset given coordinates in a permuted shared memory.
|
||||
* \param i The row index.
|
||||
* \param j The column index.
|
||||
*/
|
||||
static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
|
||||
if constexpr (swizzle_mode == SwizzleMode::k128B) {
|
||||
return i * stride + (j ^ (i % 8));
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
|
||||
return i * stride + (j ^ ((i / 2) % 4));
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k32B) {
|
||||
return i * stride + j;
|
||||
}
|
||||
}
|
||||
|
||||
/*!
|
||||
* \tparam step_size The step size to advance the offset in the permuted shared memory.
|
||||
* \param offset The current offset.
|
||||
*/
|
||||
template <uint32_t step_size>
|
||||
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
|
||||
if constexpr (swizzle_mode == SwizzleMode::k128B) {
|
||||
static_assert(step_size % 8 == 0,
|
||||
"Unsupported step size");
|
||||
return offset + step_size;
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
|
||||
static_assert(step_size == 4, "Unsupported step size");
|
||||
return offset + step_size;
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k32B) {
|
||||
static_assert(step_size == 2, "Unsupported step size");
|
||||
return offset + step_size;
|
||||
}
|
||||
}
|
||||
|
||||
// ! use with care
|
||||
template <uint32_t step_size>
|
||||
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
|
||||
if constexpr (swizzle_mode == SwizzleMode::k128B) {
|
||||
static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
|
||||
"Unsupported step size");
|
||||
if constexpr (step_size == 2) {
|
||||
return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
|
||||
} else if constexpr (step_size == 4) {
|
||||
return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
|
||||
} else {
|
||||
// step_size % 8 == 0
|
||||
return offset + step_size;
|
||||
}
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
|
||||
static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
|
||||
if constexpr (step_size == 2) {
|
||||
return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
|
||||
} else {
|
||||
return offset + step_size;
|
||||
}
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k32B) {
|
||||
return offset + step_size;
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t step_size>
|
||||
static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
|
||||
if constexpr (swizzle_mode == SwizzleMode::k128B) {
|
||||
static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
|
||||
if constexpr (step_size == 4) {
|
||||
return (offset ^ 0x4) + step_size * stride;
|
||||
} else {
|
||||
// step_size % 8 == 0
|
||||
return offset + step_size * stride;
|
||||
}
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
|
||||
static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
|
||||
if constexpr (step_size == 4) {
|
||||
return (offset ^ 0x2) + step_size * stride;
|
||||
} else {
|
||||
// step_size % 8 == 0
|
||||
return offset + step_size * stride;
|
||||
}
|
||||
} else if constexpr (swizzle_mode == SwizzleMode::k32B) {
|
||||
return offset + step_size * stride;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
|
||||
b128_t* smem_ptr = base + offset;
|
||||
mma::ldmatrix_m8n8x2(R, smem_ptr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
|
||||
b128_t* smem_ptr = base + offset;
|
||||
mma::ldmatrix_m8n8x4(R, smem_ptr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
|
||||
b128_t* smem_ptr = base + offset;
|
||||
mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
|
||||
}
|
||||
|
||||
template <cp_async::SharedMemFillMode fill_mode, typename T>
|
||||
__device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
|
||||
b128_t* smem_ptr = base + offset;
|
||||
cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
|
||||
smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
|
||||
b128_t* smem_ptr = base + offset;
|
||||
cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
|
||||
*reinterpret_cast<b128_t*>(gptr) = *(base + offset);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sageattention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,409 @@
|
||||
// Sol-Attn diagonal-threshold algorithm:
|
||||
// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
|
||||
#include "core/ggml_extensions.h"
|
||||
#include "runtime.h"
|
||||
|
||||
#if __CUDACC_VER_MAJOR__ >= 12
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <algorithm>
|
||||
#include <climits>
|
||||
|
||||
namespace sd_cuda {
|
||||
|
||||
static __device__ __forceinline__ float warp_reduce_sum(float x) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
x += __shfl_xor_sync(0xffffffff, x, offset);
|
||||
return x;
|
||||
}
|
||||
|
||||
static constexpr int SOL_DIM = 128;
|
||||
static constexpr int SOL_BLOCK = 64;
|
||||
|
||||
struct SolWorkspace {
|
||||
size_t q, k, v, kc, vc, qm, stats, threshold, size;
|
||||
int tokens, blocks;
|
||||
|
||||
explicit SolWorkspace(const ggml_tensor* dst) {
|
||||
const auto* input = dst->src[0];
|
||||
const size_t heads = input->ne[2] * input->ne[3];
|
||||
tokens = GGML_PAD(input->ne[1], SOL_BLOCK);
|
||||
blocks = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
|
||||
size = 0;
|
||||
q = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
||||
k = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
||||
v = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
|
||||
kc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
|
||||
vc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
|
||||
qm = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
|
||||
stats = reserve(heads * 2 * SOL_DIM * sizeof(float));
|
||||
threshold = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
|
||||
}
|
||||
|
||||
size_t reserve(size_t bytes) {
|
||||
const size_t offset = size;
|
||||
size += GGML_PAD(bytes, 256);
|
||||
return offset;
|
||||
}
|
||||
};
|
||||
|
||||
static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
|
||||
const int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (index >= padded * SOL_DIM) {
|
||||
return;
|
||||
}
|
||||
const int token = index / SOL_DIM;
|
||||
const int channel = index % SOL_DIM;
|
||||
const size_t src = size_t(blockIdx.y) * length * SOL_DIM + index;
|
||||
const size_t dst = size_t(blockIdx.y) * padded * SOL_DIM + index;
|
||||
qb[dst] = __float2bfloat16(token < length ? q[src] : 0.f);
|
||||
kb[dst] = __float2bfloat16(token < length ? k[src] : 0.f);
|
||||
vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
|
||||
__float2bfloat16(token < length ? v[src] : 0.f);
|
||||
}
|
||||
|
||||
static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
|
||||
const int channel = threadIdx.x;
|
||||
const int block = blockIdx.x;
|
||||
const int count = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
|
||||
float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
|
||||
for (int i = 0; i < count; ++i) {
|
||||
const int token = block * SOL_BLOCK + i;
|
||||
const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
|
||||
sum_q += __bfloat162float(q[index]);
|
||||
sum_k += __bfloat162float(k[index]);
|
||||
sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
|
||||
}
|
||||
kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
|
||||
__float2bfloat16(count ? sum_k / count : 0.f);
|
||||
vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
|
||||
if (count) {
|
||||
qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
|
||||
const int channel = threadIdx.x;
|
||||
float sum = 0.f, sum_sq = 0.f;
|
||||
for (int i = 0; i < blocks; ++i) {
|
||||
const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
|
||||
sum += value;
|
||||
sum_sq += value * value;
|
||||
}
|
||||
const float mean = sum / blocks;
|
||||
stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel] = mean;
|
||||
stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
|
||||
}
|
||||
|
||||
static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
|
||||
const int channel = threadIdx.x;
|
||||
const float q = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
|
||||
float mean = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
|
||||
float var = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
|
||||
mean = warp_reduce_sum(mean);
|
||||
var = warp_reduce_sum(var);
|
||||
__shared__ float means[4], vars[4];
|
||||
if (channel % 32 == 0) {
|
||||
means[channel / 32] = mean;
|
||||
vars[channel / 32] = var;
|
||||
}
|
||||
__syncthreads();
|
||||
if (channel == 0) {
|
||||
mean = means[0] + means[1] + means[2] + means[3];
|
||||
var = vars[0] + vars[1] + vars[2] + vars[3];
|
||||
threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
|
||||
mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
|
||||
}
|
||||
}
|
||||
|
||||
struct SolQTile {
|
||||
uint32_t x[4];
|
||||
};
|
||||
struct SolKTile {
|
||||
uint32_t x[2];
|
||||
};
|
||||
struct SolTile {
|
||||
float x[4] = {};
|
||||
static __device__ __forceinline__ int get_i(int i) {
|
||||
return (i / 2) * 8 + threadIdx.x / 4;
|
||||
}
|
||||
static __device__ __forceinline__ int get_j(int i) {
|
||||
return (threadIdx.x % 4) * 2 + i % 2;
|
||||
}
|
||||
};
|
||||
|
||||
static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
|
||||
const auto* data = reinterpret_cast<const uint32_t*>(input);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
|
||||
}
|
||||
static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
|
||||
const auto* data = reinterpret_cast<const uint32_t*>(input);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 2; ++i)
|
||||
tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
|
||||
}
|
||||
static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
auto* dst = reinterpret_cast<uint32_t*>(output.x);
|
||||
asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
|
||||
: "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
|
||||
: "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
|
||||
#else
|
||||
asm("trap;");
|
||||
#endif
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 8; ++n) {
|
||||
scores[n] = SolTile{};
|
||||
#pragma unroll
|
||||
for (int d = 0; d < 8; ++d) {
|
||||
SolKTile kt;
|
||||
sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
|
||||
sol_mma(scores[n], q[d], kt);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
scores[n].x[i] *= scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <bool PROXY>
|
||||
static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
|
||||
float maximum[2] = {-INFINITY, -INFINITY};
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 8; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const int column = n * 8 + SolTile::get_j(i);
|
||||
const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
|
||||
: (start + column < length);
|
||||
scores[n].x[i] = valid ? scores[n].x[i] : -INFINITY;
|
||||
maximum[i / 2] = fmaxf(maximum[i / 2], scores[n].x[i]);
|
||||
}
|
||||
}
|
||||
float alpha[2], sums[2] = {0.f, 0.f};
|
||||
#pragma unroll
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
|
||||
maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
|
||||
const float next = fmaxf(row_max[row], maximum[row]);
|
||||
alpha[row] = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
|
||||
row_max[row] = next;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 8; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const int column = n * 8 + SolTile::get_j(i);
|
||||
const float p = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
|
||||
const int count = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
|
||||
sums[i / 2] += p * count;
|
||||
probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
|
||||
sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
|
||||
row_sum[row] = row_sum[row] * alpha[row] + sums[row];
|
||||
}
|
||||
__syncwarp();
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 16; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
output[n].x[i] *= alpha[i / 2];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int d = 0; d < 4; ++d) {
|
||||
SolQTile pt;
|
||||
SolKTile vt;
|
||||
sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
|
||||
sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
|
||||
sol_mma(output[n], pt, vt);
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
static __global__ __launch_bounds__(128) void sol_forward(
|
||||
const nv_bfloat16* q,
|
||||
const nv_bfloat16* k,
|
||||
const nv_bfloat16* v,
|
||||
const nv_bfloat16* kc,
|
||||
const nv_bfloat16* vc,
|
||||
const float* thresholds,
|
||||
float* dst,
|
||||
int length,
|
||||
int padded,
|
||||
int padded_blocks,
|
||||
int heads,
|
||||
float scale) {
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
const int warp = threadIdx.y;
|
||||
const int lane = threadIdx.x;
|
||||
const int tid = warp * 32 + lane;
|
||||
const int blocks = padded / SOL_BLOCK;
|
||||
const int q_start = blockIdx.x * SOL_BLOCK;
|
||||
const int q_len = min(SOL_BLOCK, length - q_start);
|
||||
const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
|
||||
SolQTile qt[8];
|
||||
#pragma unroll
|
||||
for (int d = 0; d < 8; ++d) {
|
||||
sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
|
||||
}
|
||||
SolTile output[16];
|
||||
float row_max[2] = {-INFINITY, -INFINITY};
|
||||
float row_sum[2] = {0.f, 0.f};
|
||||
__shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
|
||||
__shared__ float routes[4][SOL_BLOCK];
|
||||
__shared__ int exact[SOL_BLOCK];
|
||||
const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
|
||||
for (int group = 0; group < blocks; group += SOL_BLOCK) {
|
||||
SolTile scores[8];
|
||||
sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 8; ++n) {
|
||||
#pragma unroll
|
||||
for (int column = 0; column < 2; ++column) {
|
||||
float sum = scores[n].x[column] + scores[n].x[column + 2];
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 4);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 8);
|
||||
sum += __shfl_xor_sync(0xffffffff, sum, 16);
|
||||
if (lane < 4) {
|
||||
routes[warp][n * 8 + 2 * lane + column] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid < SOL_BLOCK) {
|
||||
const int block = group + tid;
|
||||
const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
|
||||
exact[tid] = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
|
||||
}
|
||||
__syncthreads();
|
||||
sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
|
||||
vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
|
||||
padded_blocks, group, length, exact);
|
||||
for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
|
||||
if (!exact[block]) {
|
||||
continue;
|
||||
}
|
||||
const int start = (group + block) * SOL_BLOCK;
|
||||
sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
|
||||
sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
|
||||
v + base + start, padded, start, length, nullptr);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
const int head = blockIdx.y % heads;
|
||||
const int batch = blockIdx.y / heads;
|
||||
#pragma unroll
|
||||
for (int n = 0; n < 16; ++n) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const int token = q_start + warp * 16 + SolTile::get_i(i);
|
||||
if (token < length) {
|
||||
const int channel = n * 8 + SolTile::get_j(i);
|
||||
dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
|
||||
__bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
asm("trap;");
|
||||
#endif
|
||||
}
|
||||
|
||||
static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
|
||||
SDSolAttentionParams params;
|
||||
if (!sd_ggml_sol_attn_describe(op, ¶ms))
|
||||
return false;
|
||||
const auto* q = op->src[0];
|
||||
if (q->ne[2] * q->ne[3] > 65535 || GGML_PAD(q->ne[1], SOL_BLOCK) > INT_MAX / SOL_DIM)
|
||||
return false;
|
||||
for (const auto* tensor : {op->src[0], op->src[1], op->src[2]})
|
||||
if (reinterpret_cast<uintptr_t>(tensor->data) & 3)
|
||||
return false;
|
||||
return !(reinterpret_cast<uintptr_t>(op->data) & 3);
|
||||
}
|
||||
|
||||
static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
|
||||
*workspace = SolWorkspace(nodes[0]).size;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static void sol_destroy(void*, void*) {}
|
||||
|
||||
static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
|
||||
ggml_ext_gpu_context context;
|
||||
if (!cuda_context(execution->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto* dst = nodes[0];
|
||||
SDSolAttentionParams params;
|
||||
if (!sd_ggml_sol_attn_describe(dst, ¶ms))
|
||||
return GGML_STATUS_FAILED;
|
||||
const auto* q = dst->src[0];
|
||||
const SolWorkspace work(dst);
|
||||
char* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
|
||||
auto* qb = reinterpret_cast<nv_bfloat16*>(data + work.q);
|
||||
auto* kb = reinterpret_cast<nv_bfloat16*>(data + work.k);
|
||||
auto* vb = reinterpret_cast<nv_bfloat16*>(data + work.v);
|
||||
auto* kc = reinterpret_cast<nv_bfloat16*>(data + work.kc);
|
||||
auto* vc = reinterpret_cast<nv_bfloat16*>(data + work.vc);
|
||||
auto* qm = reinterpret_cast<float*>(data + work.qm);
|
||||
auto* stats = reinterpret_cast<float*>(data + work.stats);
|
||||
auto* threshold = reinterpret_cast<float*>(data + work.threshold);
|
||||
const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
|
||||
const int blocks = work.tokens / SOL_BLOCK;
|
||||
const float scale = params.scale * 1.4426950408889634f;
|
||||
const auto stream = static_cast<cudaStream_t>(context.stream);
|
||||
sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
|
||||
static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
|
||||
qb, kb, vb, length, work.tokens);
|
||||
sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
|
||||
sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
|
||||
sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
|
||||
sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
|
||||
static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
|
||||
return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
|
||||
void register_sol_attention(ggml_backend_dev_t device) {
|
||||
const auto info = cuda_device_info(device);
|
||||
if (info.device < 0)
|
||||
return;
|
||||
int previous;
|
||||
cudaDeviceProp properties;
|
||||
cudaFuncAttributes attributes;
|
||||
if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
|
||||
cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
|
||||
return;
|
||||
const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
|
||||
if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
|
||||
std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
|
||||
return;
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.sol_attn.cuda.v1";
|
||||
desc.op = GGML_OP_CUSTOM;
|
||||
desc.custom_op = sd_ggml_sol_attn_op();
|
||||
desc.priority = 100;
|
||||
desc.supports = sol_supports;
|
||||
desc.prepare = sol_prepare;
|
||||
desc.execute = sol_execute;
|
||||
desc.destroy = sol_destroy;
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
|
||||
}
|
||||
#else
|
||||
namespace sd_cuda {
|
||||
void register_sol_attention(ggml_backend_dev_t) {}
|
||||
}
|
||||
#endif
|
||||
Reference in new issue
Block a user