diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index 49e77bcef..f1e921aa7 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -123,13 +123,18 @@ Both formats imply convrot rotation even without a `convrot` flag; missing Both formats reuse the convrot activation cache when the input scale is `1` and produce F32 outputs with optional bias. CPU provides a portable implementation. -CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS; -CUDA and supported AMD targets use packed byte-dot instructions, with portable -byte arithmetic for other targets. Vulkan requires integer dot product support -and at least 13312 bytes of shared memory for W4A4 or 16452 bytes for W4A8. -GPU W4 kernels require the logical input width to be divisible by 64; H64 and -H256 are supported. Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA -by the active compilation target's +CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS. +With CUTLASS enabled, NVIDIA +SM80 and newer devices use mixed-input S8 x S4 GEMM for eligible W4A4 shapes, +reusing the INT8 scale epilogue without expanding the full weight tensor. W4A8 +uses CUTLASS integer MMA per 16-value group with FP32 scale restoration. +Other CUDA shapes and builds retain the packed byte-dot path. Set +`GGML_CUDA_W4_FORCE_DP4A=1` to compare against that path. Supported AMD targets +use packed byte-dot instructions, with portable byte arithmetic for other targets. +Vulkan requires integer dot product support and at least 13312 bytes of shared +memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input +width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU +fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout. There is no GPU architecture allowlist: newer targets supporting either intrinsic automatically use its matching layout. Selection happens per device diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 593445bf0..e2dd3f92c 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -27,6 +27,12 @@ if(SD_CUDA) 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}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include") + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr) + if(MSVC) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + endif() set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS) 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) diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index 69f91d2ae..9d4d92e1c 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -143,8 +143,9 @@ namespace sd_i8_cutlass { dst[i] = src[i / 2 + (i % 2) * (n / 2)]; } - template + template struct DenseTypes { + using Weight = std::conditional_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; using Swizzle = std::conditional_t>; @@ -161,10 +162,10 @@ namespace sd_i8_cutlass { cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, - int8_t, + Weight, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, - 16, + PackedW4 ? 32 : 16, float, cutlass::layout::RowMajor, 4, @@ -178,11 +179,12 @@ namespace sd_i8_cutlass { Epilogue, Swizzle, 4, - cutlass::arch::OpMultiplyAddSaturate>::GemmKernel; + std::conditional_t>::GemmKernel; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; template struct SwigluTypes { + using Weight = int8_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; using Swizzle = std::conditional_t>; @@ -205,6 +207,7 @@ namespace sd_i8_cutlass { DeviceInfo device; int64_t rows, n, k; bool swiglu; + bool packed_w4; 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; @@ -213,7 +216,7 @@ namespace sd_i8_cutlass { 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}; + return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast(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) { @@ -272,6 +275,8 @@ namespace sd_i8_cutlass { } template static bool configure_tile(Plan& p) { + if (p.packed_w4) + return configure_kernel>(p); if (p.swiglu) return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); @@ -300,20 +305,21 @@ namespace sd_i8_cutlass { } } 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) { +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, bool packed_w4) { #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) || + if ((packed_w4 && (swiglu || k > INT_MAX / 1024)) || 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; + auto* p = new SDI8CudaGemmPlan; + p->device = info; + p->rows = rows; + p->n = n; + p->k = k; + p->swiglu = swiglu; + p->packed_w4 = packed_w4; if (!sd_i8_cutlass::configure(*p)) { delete p; return nullptr; diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h index 189df6ef0..5e4ac9f8c 100644 --- a/src/core/cuda/i8_gemm.h +++ b/src/core/cuda/i8_gemm.h @@ -3,7 +3,7 @@ #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); +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, bool packed_w4 = false); 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/w4.cu b/src/core/cuda/w4.cu index ef46d0236..98f810e13 100644 --- a/src/core/cuda/w4.cu +++ b/src/core/cuda/w4.cu @@ -1,6 +1,13 @@ #include #include "core/ggml_extensions.h" #include "runtime.h" +#ifdef SD_CUDA_W4_CUTLASS +#include +#include +#include +#include "i8_gemm.h" +#include "w4_mma.h" +#endif namespace SD_CUDA_NAMESPACE { namespace sd_w4_cuda { @@ -789,7 +796,37 @@ namespace SD_CUDA_NAMESPACE { return false; return true; } - static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) { + struct Plan { + bool mma = false; + size_t gemm_offset = 0; +#ifdef SD_CUDA_W4_CUTLASS + SDI8CudaGemmPlan* gemm = nullptr; + ~Plan() { sd_cuda_i8_gemm_free(gemm); } +#endif + }; + static bool use_mma(ggml_backend_t backend, const SDW4Operation& op) { +#ifdef SD_CUDA_W4_CUTLASS + if (!aligned(op.input, 16) || !aligned(op.weight, 16)) + return false; + static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_W4_FORCE_DP4A"); return v && std::atoi(v); }(); + if (disabled) + return false; + static std::mutex mutex; + static std::unordered_map devices; + const auto device = cuda_device_info(ggml_backend_get_device(backend)).device; + std::lock_guard lock(mutex); + auto found = devices.find(device); + if (found != devices.end()) + return found->second; + cudaDeviceProp prop = {}; + const bool available = cudaGetDeviceProperties(&prop, device) == cudaSuccess && prop.major >= 8; + devices.emplace(device, available); + return available; +#else + return false; +#endif + } + static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int, void** state, size_t* bytes, void*) { SDW4Operation op; if (!sd_ggml_w4_describe(nodes[0], &op)) return GGML_STATUS_FAILED; @@ -812,7 +849,22 @@ namespace SD_CUDA_NAMESPACE { may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst)); } } - *bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 ? 1 : 0) + (may_alias ? 1 : 0)); + auto* plan = new Plan; + plan->mma = use_mma(backend, op); + *state = plan; + size_t gemm_bytes = 0; +#ifdef SD_CUDA_W4_CUTLASS + if (plan->mma && op.kind == SDW4Kind::W4A4) { + if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16)) + plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true); + plan->mma = plan->gemm != nullptr; + if (plan->gemm && gemm_bytes == 0) + gemm_bytes = 256; + } +#endif + const size_t tensor_bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 && !plan->mma ? 1 : 0) + (may_alias ? 1 : 0)); + plan->gemm_offset = (tensor_bytes + 255) / 256 * 256; + *bytes = plan->gemm_offset + gemm_bytes; return GGML_STATUS_SUCCESS; } static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) { @@ -821,7 +873,7 @@ namespace SD_CUDA_NAMESPACE { const auto a = reinterpret_cast(dst->data), b = reinterpret_cast(src->data); return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src); } - static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) { + static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void* state, void*) { auto* dst = nodes[0]; SDW4Operation op; ggml_ext_gpu_context context; @@ -840,7 +892,21 @@ namespace SD_CUDA_NAMESPACE { const auto* ws = static_cast(op.weight_scale->data); const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16); - if (op.kind == SDW4Kind::W4A4) { +#ifdef SD_CUDA_W4_CUTLASS + auto* plan = static_cast(state); + if (plan->mma) { + if (plan->gemm) { + auto* workspace = scratch ? reinterpret_cast(scratch) + plan->gemm_offset : nullptr; + if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast(x), reinterpret_cast(w), xs, ws, output, workspace, stream) != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + } else { + const dim3 tiles(unsigned((n + 63) / 64), unsigned((rows + 127) / 128)); + sd_w4_mma::w4a8<128, 64, 128><<>>(reinterpret_cast(w), reinterpret_cast(x), xs, ws, + static_cast(op.s_channel->data), static_cast(op.s_rel->data), bias, output, int(rows), int(n), int(k)); + } + } else +#endif + if (op.kind == SDW4Kind::W4A4) { auto* accum = reinterpret_cast(scratch + (alias ? ggml_nelements(dst) : 0)); mul_mat_w4a4_convrot_dp4a<<>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2)); if (cuda_last_error() != cuda_success) @@ -867,7 +933,7 @@ namespace SD_CUDA_NAMESPACE { d.supports = sd_w4_cuda::supports; d.prepare = sd_w4_cuda::prepare; d.execute = sd_w4_cuda::execute; - d.destroy = [](void*, void*) {}; + d.destroy = [](void* state, void*) { delete static_cast(state); }; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } } diff --git a/src/core/cuda/w4_mma.h b/src/core/cuda/w4_mma.h new file mode 100644 index 000000000..94629bfbd --- /dev/null +++ b/src/core/cuda/w4_mma.h @@ -0,0 +1,138 @@ +#ifndef __SD_CORE_CUDA_W4_MMA_H__ +#define __SD_CORE_CUDA_W4_MMA_H__ + +#include +#include +#include + +namespace sd_w4_mma { + static __device__ __forceinline__ float fp8_e4m3(uint8_t v) { + const int exponent = (v >> 3) & 15, mantissa = v & 7; + const float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9) + : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7); + return v & 128 ? -result : result; + } + + template + static __global__ void __launch_bounds__(256) w4a8( + const uint8_t* __restrict__ w, + const int8_t* __restrict__ x, + const float* __restrict__ xs, + const float* __restrict__ codebook, + const float* __restrict__ channel, + const uint8_t* __restrict__ relative, + const float* __restrict__ bias, + float* __restrict__ dst, + int rows, + int n, + int k) { + const int row_base = blockIdx.y * M, col_base = blockIdx.x * N; + const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32; + __shared__ int cb[16]; + __shared__ float cbscale; + if (tid == 0) { + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 16; ++i) + amax = fmaxf(amax, fabsf(codebook[i])); + cbscale = amax / 127.f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int value = cbscale == 0.f ? 0 : __float2int_rn(codebook[i] / cbscale); + cb[i] = max(-127, min(127, value)); + } + } + __syncthreads(); +#if __CUDA_ARCH__ >= 800 + constexpr int stride = K / 4 + 4; + __shared__ __align__(16) uint32_t a[M][stride], b[N][stride]; + __shared__ uint32_t lut[256]; + __shared__ float scales[K / 16][N]; + lut[tid] = (cb[tid & 15] & 255) | ((cb[tid >> 4] & 255) << 8); + __syncthreads(); + constexpr int warp_cols = N / 32; + constexpr int row_tiles = M / (16 * (8 / warp_cols)); + const int wr = (warp / warp_cols) * 16 * row_tiles, wc = (warp % warp_cols) * 32; + float acc[row_tiles][4][4] = {}; + for (int base = 0; base < k; base += K) { + for (int i = tid; i < M * K / 16; i += 256) { + const int r = i / (K / 16), t = (i % (K / 16)) * 16; + const int4 value = row_base + r < rows && base + t < k ? *reinterpret_cast(x + int64_t(row_base + r) * k + base + t) : make_int4(0, 0, 0, 0); + *reinterpret_cast(&a[r][t / 4]) = value; + } + for (int i = tid; i < N * K / 8; i += 256) { + const int c = i / (K / 8), t = i % (K / 8); + const uint32_t value = col_base + c < n && base + t * 8 < k ? *reinterpret_cast(w + int64_t(col_base + c) * (k / 2) + base / 2 + t * 4) : 0; + b[c][2 * t] = lut[value & 255] | (lut[(value >> 8) & 255] << 16); + b[c][2 * t + 1] = lut[(value >> 16) & 255] | (lut[value >> 24] << 16); + } + for (int i = tid; i < N * K / 16; i += 256) { + const int c = i % N, g = i / N; + scales[g][c] = col_base + c < n && base + g * 16 < k ? cbscale * (channel[col_base + c] * fp8_e4m3(relative[int64_t(col_base + c) * (k / 16) + base / 16 + g])) : 0.f; + } + __syncthreads(); +#pragma unroll + for (int g = 0; g < K / 16; ++g) { +#pragma unroll + for (int r = 0; r < row_tiles; ++r) { + const int ar = wr + r * 16 + lane / 4, ac = g * 4 + lane % 4; + const uint32_t a0 = a[ar][ac], a1 = a[ar + 8][ac]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t b0 = b[wc + j * 8 + lane / 4][ac]; + using Mma = cutlass::arch::Mma, 32, + int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor, + int, cutlass::layout::RowMajor, cutlass::arch::OpMultiplyAddSaturate>; + Mma::FragmentA af; + Mma::FragmentB bf; + Mma::FragmentC d; + reinterpret_cast(&af)[0] = a0; + reinterpret_cast(&af)[1] = a1; + reinterpret_cast(&bf)[0] = b0; + d.clear(); + // Each K=16 MMA has its own scale; combining groups before + // dequantization would change the W4A8 representation. + Mma{}(d, af, bf, d); + const float s0 = scales[g][wc + j * 8 + (lane % 4) * 2], s1 = scales[g][wc + j * 8 + (lane % 4) * 2 + 1]; +#pragma unroll + for (int v = 0; v < 4; ++v) + acc[r][j][v] = fmaf(float(d[v]), v % 2 ? s1 : s0, acc[r][j][v]); + } + } + } + __syncthreads(); + } +#pragma unroll + for (int r = 0; r < row_tiles; ++r) { +#pragma unroll + for (int j = 0; j < 4; ++j) { +#pragma unroll + for (int v = 0; v < 4; ++v) { + const int row = row_base + wr + r * 16 + lane / 4 + (v / 2) * 8; + const int col = col_base + wc + j * 8 + (lane % 4) * 2 + v % 2; + if (row < rows && col < n) + dst[int64_t(row) * n + col] = __fmul_rn(acc[r][j][v], xs[row]) + (bias ? bias[col] : 0.f); + } + } + } +#else + for (int i = tid; i < M * N; i += 256) { + const int row = row_base + i / N, col = col_base + i % N; + if (row < rows && col < n) { + float sum = 0.f; + for (int g = 0; g < k / 16; ++g) { + int d = 0; + for (int j = 0; j < 16; ++j) { + const int t = g * 16 + j; + d += cb[(w[int64_t(col) * (k / 2) + t / 2] >> (4 * (t % 2))) & 15] * int(x[int64_t(row) * k + t]); + } + sum = fmaf(float(d), cbscale * (channel[col] * fp8_e4m3(relative[int64_t(col) * (k / 16) + g])), sum); + } + dst[int64_t(row) * n + col] = __fmul_rn(sum, xs[row]) + (bias ? bias[col] : 0.f); + } + } +#endif + } + +} +#endif