diff --git a/CMakeLists.txt b/CMakeLists.txt index d89fe34d..b5e5208c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -348,17 +348,18 @@ if (SD_CUDA) 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 + target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu) + set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu PROPERTIES INCLUDE_DIRECTORIES "${GGML_CUDA_CUTLASS_PATH}/include" COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr") if (MSVC) - set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) endif() if (GGML_CUDA_FORCE_CUBLAS) - set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) + set_property(SOURCE src/core/cuda/i8_projection.cu src/core/cuda/i8_gated_residual.cu src/core/cuda/i8_gemm.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS) endif() set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION) + set_property(SOURCE src/core/cuda/int8.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS) endif() set_target_properties(${SD_LIB} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON) target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 4077aebf..408d3622 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -191,11 +191,33 @@ ConvRot, and rescaling kernels. F32 inputs are quantized into workspace; ordinar I8 inputs retain separate scale tensors and use a padded temporary only when needed by cuBLAS. Packed ConvRot inputs are consumed directly. Large ConvRot rows use multiple kernels when they exceed shared memory. Aliasing outputs use scratch. -These registrations do not require CUTLASS and do not intercept builtin nodes yet. +The registrations also intercept builtin INT8 GEMM and standalone ConvRot nodes. CUDA GEMM requires Turing or newer, contiguous tensors, K/N divisible by four, -float-aligned inputs, and at most 65535 rows; CUDA ConvRot uses groups of 256. +float-aligned inputs, and row counts that fit a padded signed 32-bit dimension; +CUDA ConvRot uses groups of 256. Unsupported custom nodes can fall back to CPU through the scheduler. +When CUTLASS is configured, `src/core/cuda/i8_gemm.cu` supplies dense and SwiGLU +INT8 GEMM kernels on Ampere or newer. It preserves the native tile selection, +small-row policy, Stream-K scheduling, Ada column bands, and separate scaling +order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row +concatenation for builtin and custom GEMM nodes. Weight scales remain explicit +inputs, including the gate/up scale permutation recorded on every SwiGLU launch. +Plans reserve bounded Stream-K workspace, cache occupancy on first execution, +and initialize barriers on every launch, so CUDA Graph replay observes updated +inputs and scales. External consumers retain ordinary execution; overlapping +outputs use workspace. The CUTLASS SwiGLU/ConvRot chain can write its packed +output directly after GEMM, because quantization reads only its workspace +intermediate. Shapes or alignments unsuitable for CUTLASS use cuBLAS. +Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to +preserve their previous accumulation behavior. Fusion disable and force-cuBLAS +settings are honored. No CUTLASS headers are needed for the cuBLAS path. + +The native CUDA CUTLASS GEMM implementation and its SwiGLU/row-concat matchers +have been removed from ggml. Native cuBLAS GEMM and standalone ConvRot remain +compatibility implementations for callers without these registrations. HIP +retains its native INT8/activation fusion paths. + `src/core/cuda/i8_gated_residual.cu` implements `INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered CUDA fusion, for both builtin GEMM and custom I8 nodes with separate activation @@ -230,11 +252,9 @@ 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. +The corresponding scale, normalization, and unary-activation fusion rules in +ggml now only serve HIP. CUDA GEMM/SwiGLU/ConvRot chains are selected by the +application INT8 extension before the standalone activation fusion. CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and `GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its @@ -278,7 +298,7 @@ CUDA VMM so native scratch addresses remain valid for the backend lifetime; legacy scratch pools retain the existing fragment path. HIP and MUSA do not yet provide this plan executor. sd.cpp's CUDA registrations opt in to reusable replay. -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. +Remaining CUDA fusion rules and Vulkan INT8/ConvRot implementations 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 92408021..1daaba22 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 924080214ed35e9a1251ee61bc3185f105c52b41 +Subproject commit 1daaba22f2459202f292298e935939a47f00cc8c diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu index 82d6554a..eba72cc1 100644 --- a/src/core/cuda/convrot.cu +++ b/src/core/cuda/convrot.cu @@ -145,11 +145,6 @@ namespace sd_convrot { 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; } diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu new file mode 100644 index 00000000..84d23326 --- /dev/null +++ b/src/core/cuda/i8_gemm.cu @@ -0,0 +1,330 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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 + struct TileConfig { + using Tile = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + }; + + 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 cache; + std::lock_guard 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_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(ld)), half_n(extent.strided() / 2) {} + CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld) + : stride(static_cast(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 + struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore> { + using Base = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Base::Base; + template + 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(&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(this->params_ptr->ptr_aux + offset) = value; + } + } + } + }; + template + CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) { + auto cb = Base::get_callbacks(tile, tid, shape); + return Callbacks(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 + struct DenseTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Store = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + 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; + }; + template + struct SwigluTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Store = SwigluStore; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal, cutlass::epilogue::thread::LinearCombination, 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; + using Kernel = std::conditional_t, + cutlass::gemm::kernel::GemmWithEpilogueVisitor>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + }; + 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 + 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 + 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(x) | reinterpret_cast(w) | reinterpret_cast(xs) | + reinterpret_cast(ws) | reinterpret_cast(out)) % + 16) + return GGML_STATUS_FAILED; + auto* paired = static_cast(scratch); + auto args = arguments(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, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess) + return GGML_STATUS_FAILED; + if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2, 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<<>>(ws, paired, int(p.n)); + if (cudaGetLastError() != cudaSuccess) + return GGML_STATUS_FAILED; + } + if (params.init_workspace(static_cast(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess) + return GGML_STATUS_FAILED; + cutlass::arch::synclog_setup(); + cutlass::Kernel2<<>>(params); + return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + template + static bool configure_kernel(Plan& p) { + using Kernel = typename Types::Kernel; + if (sizeof(typename Kernel::SharedStorage) > p.device.shared) + return false; + auto args = arguments(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; + return true; + } + template + static bool configure_tile(Plan& p) { + if (p.swiglu) + return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); + return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); + } + static bool configure(Plan& p) { + auto try_tile = [&](Tile tile) { + switch (tile) { + case TILE_128_256: + return configure_tile>(p); + case TILE_128_128: + return configure_tile>(p); + case TILE_64_128: + return configure_tile>(p); + case TILE_16_128: + return configure_tile>(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; +} diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h new file mode 100644 index 00000000..189df6ef --- /dev/null +++ b/src/core/cuda/i8_gemm.h @@ -0,0 +1,9 @@ +#ifndef __SD_CORE_CUDA_I8_GEMM_H__ +#define __SD_CORE_CUDA_I8_GEMM_H__ +#include +#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 diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index e1ca6cba..b0abd860 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -1,4 +1,7 @@ +#include +#include "convrot_kernel.h" #include "core/int8_ops.h" +#include "i8_gemm.h" #include #include @@ -301,7 +304,7 @@ namespace sd_i8_cuda { int64_t n, int64_t rows) { const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; - const int64_t row = blockIdx.y; + 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]; @@ -322,14 +325,16 @@ namespace sd_i8_cuda { const float4* weight_scales, const float4* bias, float4* dst, - int64_t n4) { + int64_t n4, + int64_t rows) { const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (output >= n4) { + const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535; + if (output >= n4 || row >= rows) { return; } - const int64_t index = (int64_t)blockIdx.y * n4 + output; + const int64_t index = row * n4 + output; const int4 v = src[index]; - const float scale = scales[blockIdx.y]; + 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]; @@ -371,7 +376,7 @@ namespace sd_i8_cuda { cache.emplace(device, info); return info; } - static bool aligned(const ggml_tensor* tensor) { + static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) { if (!tensor) return true; size_t offset = 0; @@ -379,14 +384,14 @@ namespace sd_i8_cuda { offset += tensor->view_offs; tensor = tensor->view_src; } - return (reinterpret_cast(tensor->data) + offset) % 4 == 0; + return (reinterpret_cast(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 > 65535 || k % 4 || n % 4 || (op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256)))) + 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)) @@ -402,41 +407,105 @@ namespace sd_i8_cuda { 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) && device_info(device).ordinal >= 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(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 <= device_info(device).shared) + return 3; + } + return 2; + } struct Plan { DeviceInfo device; + Kind kind; int64_t k, n, rows, padded; - size_t quant, scales, accum, partial, output; + 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, void** state, size_t* workspace, void*) { - const auto* dst = nodes[0]; + 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(dst, &op); + const bool mat = sd_ggml_i8_describe(nodes[0], &op); auto* p = new Plan{}; *state = p; - p->device = device_info(ggml_backend_get_device(backend)); + auto device = ggml_backend_get_device(backend); + p->device = 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 ? dst->ne[0] : 0; - p->rows = mat ? ggml_nrows(dst) : ggml_nrows(dst->src[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); - p->accum = reserve(total, size_t(p->n) * p->padded * sizeof(int32_t)); +#ifdef SD_CUDA_I8_CUTLASS + size_t gemm_bytes = 0; + // Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow. + if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || 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; p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); - p->output = reserve(total, ggml_nbytes(dst)); + 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*) { - delete static_cast(state); + auto* p = static_cast(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) @@ -474,7 +543,54 @@ namespace sd_i8_cuda { } return cudaGetLastError(); } - static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) { + template + 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(src)[row * n + col]) * xs[row]; + up = float(reinterpret_cast(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 + static cudaError_t quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, cudaStream_t stream) { + const int64_t hidden = p.n / 2; + auto kernel = sd_convrot::quantize; + static std::mutex mutex; + static std::unordered_map initialized; + { + std::lock_guard lock(mutex); + if (!initialized[p.device.ordinal]) { + auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128)); + if (error != cudaSuccess) + return error; + initialized[p.device.ordinal] = true; + } + } + if (p.padded > p.rows) { + auto error = cudaMemsetAsync(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream); + if (error != cudaSuccess) + return error; + } + const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256 + : 1024) + : 256; + kernel<<>>(src, output, + reinterpret_cast(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 cudaGetLastError(); + } + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { const auto& p = *static_cast(state); 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")); @@ -483,12 +599,21 @@ namespace sd_i8_cuda { return GGML_STATUS_FAILED; auto stream = static_cast(context.stream); auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); - auto* dst = nodes[0]; + auto* dst = nodes[count - 1]; SDI8Operation op; - const bool mat = sd_ggml_i8_describe(dst, &op); + const bool mat = sd_ggml_i8_describe(nodes[0], &op); bool alias = false; - for (auto* input : dst->src) - alias |= overlaps(dst, input); + 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_t error = cudaSuccess; if (!mat) { @@ -496,10 +621,6 @@ namespace sd_i8_cuda { auto* scales = reinterpret_cast(quant + p.k * p.padded); error = convrot_launch(p, stream, static_cast(dst->src[0]->data), quant, scales, reinterpret_cast(scratch + p.partial)); } else { - auto blas_query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1")); - auto handle = blas_query ? static_cast(blas_query(execution->backend)) : nullptr; - if (!handle) - return GGML_STATUS_FAILED; auto* quant = reinterpret_cast(scratch + p.quant); const float* scales; if (op.input->type == GGML_TYPE_F32) { @@ -526,26 +647,64 @@ namespace sd_i8_cuda { } if (error != cudaSuccess || cudaGetLastError() != cudaSuccess) return GGML_STATUS_FAILED; - auto* accum = reinterpret_cast(scratch + p.accum); - const int32_t alpha = 1, beta = 0; - if (cublasGemmEx(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) != CUBLAS_STATUS_SUCCESS) - return GGML_STATUS_FAILED; - const auto* ws = op.weight_scale ? static_cast(op.weight_scale->data) : nullptr; - const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; - if ((reinterpret_cast(ws) | reinterpret_cast(bias) | reinterpret_cast(output)) % 16 == 0) { - const dim3 grid(unsigned((p.n + 1023) / 1024), unsigned(p.rows)); - dequantize_i32_rows_vec4_cuda<<>>(reinterpret_cast(accum), scales, reinterpret_cast(ws), reinterpret_cast(bias), static_cast(output), p.n / 4); - } else { - const dim3 grid(unsigned((p.n + 255) / 256), unsigned(p.rows)); - dequantize_i32_rows_cuda<<>>(accum, scales, ws, bias, static_cast(output), p.n, p.rows); + const auto* ws = op.weight_scale ? static_cast(op.weight_scale->data) : nullptr; + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0; + auto* gemm_output = reinterpret_cast(static_cast(output) + prefix); + auto* accum = reinterpret_cast(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(op.weight->data), scales, ws, + intermediate ? reinterpret_cast(accum) : gemm_output, scratch + p.gemm_offset, stream); + if (status != GGML_STATUS_SUCCESS) + return status; + } else +#endif + { + auto blas_query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1")); + auto handle = blas_query ? static_cast(blas_query(execution->backend)) : nullptr; + if (!handle) + return GGML_STATUS_FAILED; + const int32_t alpha = 1, beta = 0; + if (cublasGemmEx(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) != CUBLAS_STATUS_SUCCESS) + return GGML_STATUS_FAILED; } - error = cudaGetLastError(); + const auto* values = reinterpret_cast(accum); + if (p.kind == SWIGLU_QUANTIZE) { + if (p.fast_swiglu) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else if (p.gemm) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else + error = quantize_epilogue(p, values, scales, ws, static_cast(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<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + else + swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + error = cudaGetLastError(); + } else if (!p.gemm) { + const unsigned gy = unsigned(std::min(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535); + if ((reinterpret_cast(ws) | reinterpret_cast(bias) | reinterpret_cast(gemm_output)) % 16 == 0) { + dequantize_i32_rows_vec4_cuda<<>>( + reinterpret_cast(accum), scales, reinterpret_cast(ws), reinterpret_cast(bias), reinterpret_cast(gemm_output), p.n / 4, p.rows); + } else { + dequantize_i32_rows_cuda<<>>(accum, scales, ws, bias, gemm_output, p.n, p.rows); + } + error = cudaGetLastError(); + } + if (error == cudaSuccess && prefix) + error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream); } 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(ggml_backend_dev_t device) { @@ -560,8 +719,27 @@ void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) { 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.cuda.override.v1"; + desc.op = GGML_OP_MUL_MAT; + desc.custom_op = nullptr; + 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.cuda.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.name = "sd.linear_i8.epilogue.cuda.custom.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_i8_op(); + 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.cuda.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)); + desc.name = "sd.quantize_i8_convrot.cuda.override.v1"; + desc.op = GGML_OP_QUANTIZE_I8_CONVROT; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); }