mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-10 19:06:33 +00:00
perf: accelerate packed INT4 CUDA inference with CUTLASS (#2278)
This commit is contained in:
1 parent
7952693004
commit
f5df5f8b84
6 files changed
+247
-26
No files matched your search
+12
-7
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
+19
-13
@@ -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;
|
||||
|
||||
@@ -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
|
||||
+71
-5
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
Reference in new issue
Block a user