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

This commit is contained in:
leejet authored and GitHub committed 2026-11-06 01:59:05 +08:00
1 parent 7952693004
commit f5df5f8b84
6 files changed
+240 -19

No files matched your search

+12 -7
View File
@@ -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
+6
View File
@@ -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)
+13 -7
View File
@@ -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,12 +305,12 @@ 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;
@@ -314,6 +319,7 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row
p->n = n;
p->k = k;
p->swiglu = swiglu;
p->packed_w4 = packed_w4;
if (!sd_i8_cutlass::configure(*p)) {
delete p;
return nullptr;
+1 -1
View File
@@ -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
+70 -4
View File
@@ -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,6 +892,20 @@ 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);
#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));
@@ -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));
}
}
+138
View File
@@ -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