perf: accelerate packed INT4 CUDA inference with CUTLASS (#2278)

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)

src/core/cuda/i8_gemm.cu

@@ -143,8 +143,9 @@ namespace sd_i8_cutlass {
             dst[i] = src[i / 2 + (i % 2) * (n / 2)];
     }
 
-    template <class Config, bool StreamK>
+    template <class Config, bool StreamK, bool PackedW4 = false>
     struct DenseTypes {
+        using Weight   = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
         using Tile     = typename Config::Tile;
         using Warp     = typename Config::Warp;
         using Swizzle  = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
@@ -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<PackedW4, cutlass::arch::OpMultiplyAddMixedInputUpcast, cutlass::arch::OpMultiplyAddSaturate>>::GemmKernel;
         using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
     template <class Config, bool StreamK>
     struct SwigluTypes {
+        using Weight          = int8_t;
         using Tile            = typename Config::Tile;
         using Warp            = typename Config::Warp;
         using Swizzle         = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
@@ -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<const typename Types::Weight*>(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) {
@@ -272,6 +275,8 @@ namespace sd_i8_cutlass {
     }
     template <class Config>
     static bool configure_tile(Plan& p) {
+        if (p.packed_w4)
+            return configure_kernel<DenseTypes<Config, false, true>>(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);
@@ -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;

src/core/cuda/i8_gemm.h

@@ -3,7 +3,7 @@
 #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);
+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

src/core/cuda/w4.cu

@@ -1,6 +1,13 @@
 #include <climits>
 #include "core/ggml_extensions.h"
 #include "runtime.h"
+#ifdef SD_CUDA_W4_CUTLASS
+#include <cstdlib>
+#include <mutex>
+#include <unordered_map>
+#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<int, bool> devices;
+            const auto device = cuda_device_info(ggml_backend_get_device(backend)).device;
+            std::lock_guard<std::mutex> 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<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(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<const float*>(op.weight_scale->data);
             const auto* bias = op.bias ? static_cast<const float*>(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<Plan*>(state);
+            if (plan->mma) {
+                if (plan->gemm) {
+                    auto* workspace = scratch ? reinterpret_cast<char*>(scratch) + plan->gemm_offset : nullptr;
+                    if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast<const int8_t*>(x), reinterpret_cast<const int8_t*>(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><<<tiles, 256, 0, stream>>>(reinterpret_cast<const uint8_t*>(w), reinterpret_cast<const int8_t*>(x), xs, ws,
+                                                                             static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data), bias, output, int(rows), int(n), int(k));
+                }
+            } else
+#endif
+                if (op.kind == SDW4Kind::W4A4) {
                 auto* accum = reinterpret_cast<int32_t*>(scratch + (alias ? ggml_nelements(dst) : 0));
                 mul_mat_w4a4_convrot_dp4a<<<grid, block, 0, stream>>>(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<sd_w4_cuda::Plan*>(state); };
         GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
     }
 }

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 <cuda_runtime.h>
+#include <cutlass/arch/mma_sm80.h>
+#include <cstdint>
+
+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 <int M, int N, int K>
+    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<const int4*>(x + int64_t(row_base + r) * k + base + t) : make_int4(0, 0, 0, 0);
+                *reinterpret_cast<int4*>(&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<const uint32_t*>(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<cutlass::gemm::GemmShape<16, 8, 16>, 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<uint32_t*>(&af)[0] = a0;
+                        reinterpret_cast<uint32_t*>(&af)[1] = a1;
+                        reinterpret_cast<uint32_t*>(&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