diff --git a/CMakeLists.txt b/CMakeLists.txt index 7461b49c..d89fe34d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -345,8 +345,8 @@ if (SD_CUDA) set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA) if (NOT SD_USE_UPSTREAM_GGML) enable_language(CUDA) - target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu) - set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") + target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu) + set_source_files_properties(src/core/cuda/rope.cu src/core/cuda/int8.cu src/core/cuda/convrot.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") if (GGML_CUDA_CUTLASS_PATH) target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu) set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu PROPERTIES diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 565dfebc..56ba9dc1 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -214,7 +214,35 @@ for an exact in-place residual without concatenation. External consumers and unsupported layouts retain ordinary operator execution. Fusion and force-cuBLAS disable settings are honored. -The builtin CUDA/Vulkan INT8 GEMM, ConvRot quantization, and remaining fusion rules +`src/core/cuda/convrot.cu` registers CUDA fusion for `Scale -> ConvRot`, +`Norm -> Mul -> ConvRot`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add -> +Add(shift) -> ConvRot`, `SwiGLU -> ConvRot`, and `SiLU/Sigmoid -> Mul -> ConvRot`. +Both builtin and custom ConvRot outputs are accepted. SwiGLU accepts packed, +swapped, or separate gate/up inputs. Independent modulation and up-branch +producers are moved ahead of their consumers through the safe graph-delay API. +Reordering keeps an owning producer together with in-place consumers, such as +Krea2's RMSNorm weight multiplication, so the alias checks can preserve their +storage relationship. +The application kernel preserves the original reduction order, thread selection, +ConvRot row padding, and separate scale tail. Prepared plans cache operation +metadata, use the borrowed CUDA stream, and reserve scheduler workspace for +overlapping outputs. Long rows or external intermediate consumers retain +ordinary execution. These fusions do not require CUTLASS and honor +`GGML_CUDA_DISABLE_FUSION`. + +The corresponding scale, normalization, and unary-activation rules in ggml now +only serve HIP. CUDA retains the native INT8 GEMM/SwiGLU epilogue chain; the +ConvRot extension leaves that chain together to preserve its existing fusion. +Standalone builtin ConvRot and its native GEMM integration are still pending +migration. + +CUDA extension kernels currently execute between native graph fragments. Native +fragments retain CUDA Graph caching, but the extension execution plan is not +captured as a whole. Moving small fusions out of native fragments can therefore +increase host submission overhead even when kernel counts and GPU work stay the +same. Whole-plan CUDA Graph capture remains pending. + +The builtin CUDA/Vulkan INT8 GEMM, standalone ConvRot quantization, and remaining fusion rules still need to be externalized. Vulkan execution services and kernels also remain pending. This implementation does not yet externalize all model-specific kernels or add dynamic dtypes. diff --git a/ggml b/ggml index a6903f4a..b34bc427 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit a6903f4a06c5d3b980964fedece61587e86c5678 +Subproject commit b34bc427955068ecd7fe080df214dc94a89e3d80 diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu new file mode 100644 index 00000000..705c1e00 --- /dev/null +++ b/src/core/cuda/convrot.cu @@ -0,0 +1,320 @@ +#include +#include +#include +#include +#include +#include +#include +#include "convrot_kernel.h" +#include "core/int8_ops.h" + +namespace sd_convrot { + struct DeviceInfo { + int ordinal = -1; + size_t shared = 0; + }; + static DeviceInfo device_info(ggml_backend_dev_t device) { + static std::mutex mutex; + static std::unordered_map cache; + std::lock_guard lock(mutex); + auto found = cache.find(device); + if (found != cache.end()) + return found->second; + DeviceInfo info; + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_cuda_device_v1")); + int ordinal = query ? query(device) : -1; + cudaDeviceProp props = {}; + if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess) + info = {ordinal, size_t(props.sharedMemPerBlockOptin)}; + cache.emplace(device, info); + return info; + } + 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 DeviceInfo& device) { + const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256 + : 128; + return device.ordinal >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared; + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + if (fusion_disabled()) + return 0; + auto* g = const_cast(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); + // Native INT8 epilogues consume GEMM and GLU together; splitting them loses that fusion. + if (node->op == GGML_OP_GLU && nodes[count - 1]->op == GGML_OP_QUANTIZE_I8_CONVROT && + node->src[0]->op == GGML_OP_MUL_MAT && node->src[0]->src[0]->type == GGML_TYPE_I8 && + node->src[0]->src[2] && !node->src[0]->src[3]) + return 0; + Operation op; + return describe(nodes, count, op) && supported(op, 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, 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 { + DeviceInfo 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 = 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(state); + } + static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; + const auto x = reinterpret_cast(a->data), y = reinterpret_cast(b->data); + return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); + } + template + static cudaError_t launch(Plan& p, const Operation& op, int8_t* output, cudaStream_t stream) { + constexpr bool warp = Scaled || RMS; + constexpr int extra = Normalize ? 384 : RMS ? 256 + : 128; + if (!p.configured) { + auto status = cudaFuncSetAttribute(quantize, + cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - extra)); + if (status != cudaSuccess) + 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(t->data) : nullptr; }; + quantize<<>>( + data(op.input) + op.gate_offset, output, reinterpret_cast(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 cudaGetLastError(); + } + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { + auto& plan = *static_cast(state); + const auto& op = plan.op; + auto* dst = nodes[count - 1]; + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend)); + auto query = reinterpret_cast(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(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(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); + const int64_t padded = (op.rows + 3) / 4 * 4; + if (padded > op.rows && cudaMemsetAsync(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cudaSuccess) + return GGML_STATUS_FAILED; + cudaError_t error; + switch (op.kind) { + case SCALE: + error = launch(plan, op, output, stream); + break; + case NORM: + error = launch(plan, op, output, stream); + break; + case RMS_MOD: + error = launch(plan, op, output, stream); + break; + case SWIGLU: + case SILU: + error = launch(plan, op, output, stream); + break; + case SIGMOID: + error = launch(plan, op, output, stream); + break; + default: + return GGML_STATUS_FAILED; + } + if (error != cudaSuccess || (alias && cudaMemcpyAsync(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)) + return GGML_STATUS_FAILED; + return GGML_STATUS_SUCCESS; + } +} + +void sd_ggml_register_cuda_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) + ".cuda.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(device, &desc)); + } +} diff --git a/src/core/cuda/convrot_kernel.h b/src/core/cuda/convrot_kernel.h new file mode 100644 index 00000000..970bf88b --- /dev/null +++ b/src/core/cuda/convrot_kernel.h @@ -0,0 +1,249 @@ +#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__ +#define __SD_CORE_CUDA_CONVROT_KERNEL_H__ +#include +#include +namespace sd_convrot { + 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, 32); + return x; + } + static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + x.x += __shfl_xor_sync(0xffffffff, x.x, offset, 32); + x.y += __shfl_xor_sync(0xffffffff, x.y, offset, 32); + } + return x; + } + template + static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int offset = Width / 2; offset > 0; offset >>= 1) + x = fmaxf(x, __shfl_xor_sync(0xffffffff, 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 + __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 = __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(); + } 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 diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index d48757d4..aa005291 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -11,6 +11,7 @@ #ifdef SD_CUDA_EXTENSIONS void sd_ggml_register_cuda_i8(ggml_backend_dev_t device); +void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device); void sd_ggml_register_cuda_rope(ggml_backend_dev_t device); #endif #ifdef SD_CUDA_CUTE_ATTENTION @@ -182,6 +183,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { if (cuda) { #ifdef SD_CUDA_EXTENSIONS sd_ggml_register_cuda_i8(device); + sd_ggml_register_cuda_convrot(device); sd_ggml_register_cuda_rope(device); #endif #ifdef SD_CUDA_CUTE_ATTENTION