mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 03:24:30 +00:00
feat: support packed INT4 convrot models (#2276)
Co-authored-by: noctrex <noctrex@gmail.com>
This commit is contained in:
leejetandnoctrex authored and GitHub committed
1 parent
58bb3c1b1c
commit
7952693004
48 files changed
+2003
-102
No files matched your search
@@ -5,6 +5,7 @@ namespace sd_cpu {
|
||||
|
||||
void register_extensions(ggml_backend_dev_t device) {
|
||||
register_i8(device);
|
||||
register_w4(device);
|
||||
register_rope(device);
|
||||
register_modulate(device);
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ namespace sd_cpu {
|
||||
void register_rope(ggml_backend_dev_t device);
|
||||
void register_modulate(ggml_backend_dev_t device);
|
||||
void register_i8(ggml_backend_dev_t device);
|
||||
void register_w4(ggml_backend_dev_t device);
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
#include "core/cpu/registration.h"
|
||||
#include "core/ggml_extensions.h"
|
||||
|
||||
#ifndef SD_USE_UPSTREAM_GGML
|
||||
#include <cstring>
|
||||
|
||||
namespace sd_cpu {
|
||||
struct W4Execution {
|
||||
SDW4Operation op;
|
||||
float* output;
|
||||
int64_t rows;
|
||||
};
|
||||
|
||||
static void multiply_w4(ggml_tensor*, int ith, int nth, void* data) {
|
||||
auto& e = *static_cast<W4Execution*>(data);
|
||||
const auto& op = e.op;
|
||||
const int64_t k = op.input->ne[0], n = op.weight->ne[1];
|
||||
const auto* w = static_cast<const uint8_t*>(op.weight->data);
|
||||
const auto* x = static_cast<const int8_t*>(op.input->data);
|
||||
const auto* xs = reinterpret_cast<const float*>(x + k * ((e.rows + 3) / 4 * 4));
|
||||
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 auto* sc = op.s_channel ? static_cast<const float*>(op.s_channel->data) : nullptr;
|
||||
const auto* sr = op.s_rel ? static_cast<const uint8_t*>(op.s_rel->data) : nullptr;
|
||||
const auto to_float = ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float;
|
||||
for (int64_t index = ith; index < n * e.rows; index += nth) {
|
||||
const int64_t row = index / n, col = index % n;
|
||||
const auto* packed = w + col * (k / 2);
|
||||
float value;
|
||||
if (op.kind == SDW4Kind::W4A4) {
|
||||
uint32_t sum = 0;
|
||||
for (int64_t i = 0; i < k; ++i) {
|
||||
const int nibble = (packed[i / 2] >> (4 * (i % 2))) & 15;
|
||||
sum += uint32_t(((nibble ^ 8) - 8) * int32_t(x[row * k + i]));
|
||||
}
|
||||
int32_t signed_sum;
|
||||
std::memcpy(&signed_sum, &sum, sizeof(sum));
|
||||
value = float(signed_sum) * xs[row] * ws[col];
|
||||
} else {
|
||||
float sum = 0.f;
|
||||
for (int64_t g = 0; g < k / 16; ++g) {
|
||||
float relative;
|
||||
to_float(sr + col * (k / 16) + g, &relative, 1);
|
||||
const float scale = sc[col] * relative;
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
const int64_t i = g * 16 + j;
|
||||
const int code = (packed[i / 2] >> (4 * (i % 2))) & 15;
|
||||
sum += (ws[code] * scale) * float(x[row * k + i]);
|
||||
}
|
||||
}
|
||||
value = sum * xs[row];
|
||||
}
|
||||
e.output[index] = value + (bias ? bias[col] : 0.f);
|
||||
}
|
||||
}
|
||||
|
||||
static bool w4_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
|
||||
SDW4Operation op;
|
||||
return sd_ggml_w4_describe(node, &op);
|
||||
}
|
||||
static ggml_status w4_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
|
||||
*bytes = ggml_nbytes(nodes[0]);
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static ggml_status w4_execute(const ggml_ext_execution_context* ctx, ggml_tensor* const* nodes, int, void*, void*) {
|
||||
auto* dst = nodes[0];
|
||||
W4Execution e = {};
|
||||
if (!sd_ggml_w4_describe(dst, &e.op))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->backend));
|
||||
auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
|
||||
if (!run)
|
||||
return GGML_STATUS_FAILED;
|
||||
e.rows = ggml_nrows(dst);
|
||||
e.output = static_cast<float*>(ggml_backend_buffer_get_base(ctx->workspace));
|
||||
auto status = run(ctx->backend, dst, multiply_w4, &e);
|
||||
if (status == GGML_STATUS_SUCCESS)
|
||||
std::memcpy(dst->data, e.output, ggml_nbytes(dst));
|
||||
return status;
|
||||
}
|
||||
void register_w4(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc d = {};
|
||||
d.abi_version = GGML_EXT_ABI_VERSION;
|
||||
d.name = "sd.linear_w4_convrot.cpu.v1";
|
||||
d.op = GGML_OP_CUSTOM;
|
||||
d.custom_op = sd_ggml_w4_op();
|
||||
d.priority = 100;
|
||||
d.supports = w4_supports;
|
||||
d.prepare = w4_prepare;
|
||||
d.execute = w4_execute;
|
||||
d.destroy = [](void*, void*) {};
|
||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW)
|
||||
|
||||
function(sd_cuda_sources VENDOR LANGUAGE)
|
||||
set(SOURCES)
|
||||
foreach(OP rope int8 convrot affine modulate bf16 rms_affine registration)
|
||||
foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration)
|
||||
set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
|
||||
string(TOUPPER "${VENDOR}" UPPER_VENDOR)
|
||||
file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
|
||||
@@ -49,6 +49,8 @@ if(SD_HIPBLAS)
|
||||
set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
|
||||
endif()
|
||||
set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
|
||||
# W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12.
|
||||
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64")
|
||||
target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
|
||||
target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
|
||||
endif()
|
||||
|
||||
+65
-5
@@ -386,12 +386,55 @@ namespace SD_CUDA_NAMESPACE {
|
||||
return info.int8_blas &&
|
||||
ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
|
||||
}
|
||||
template <bool Quantize>
|
||||
static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) {
|
||||
const int64_t group = blockIdx.x, groups = k / 64, row = group / groups;
|
||||
const int tid = threadIdx.x;
|
||||
__shared__ float values[64];
|
||||
__shared__ float maxima[16];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f);
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int stride = 1; stride < 64; stride *= 4) {
|
||||
const int i = (tid / stride) * 4 * stride + tid % stride;
|
||||
const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
|
||||
values[i] = convrot_h4_cuda(a, b, c, d, 0);
|
||||
values[i + stride] = convrot_h4_cuda(a, b, c, d, 1);
|
||||
values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2);
|
||||
values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3);
|
||||
__syncthreads();
|
||||
}
|
||||
if constexpr (Quantize) {
|
||||
const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
const int value = __float2int_rn(values[tid + j * 16] * inverse);
|
||||
dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value)));
|
||||
}
|
||||
} else {
|
||||
float amax = 0.f;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
amax = fmaxf(amax, fabsf(values[tid + j * 16]));
|
||||
maxima[tid] = amax;
|
||||
__syncthreads();
|
||||
for (int stride = 8; stride; stride >>= 1) {
|
||||
if (tid < stride)
|
||||
maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
|
||||
__syncthreads();
|
||||
}
|
||||
if (!tid)
|
||||
partial[group] = maxima[0];
|
||||
}
|
||||
}
|
||||
static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
|
||||
int group;
|
||||
if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
|
||||
if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0]))
|
||||
return false;
|
||||
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) && cuda_device_info(device).device >= 0;
|
||||
return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0;
|
||||
}
|
||||
enum Kind { MAT,
|
||||
QUANTIZE,
|
||||
@@ -421,7 +464,8 @@ namespace SD_CUDA_NAMESPACE {
|
||||
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) &&
|
||||
int group;
|
||||
if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) &&
|
||||
size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
|
||||
return 3;
|
||||
}
|
||||
@@ -430,6 +474,7 @@ namespace SD_CUDA_NAMESPACE {
|
||||
struct Plan {
|
||||
ggml_ext_gpu_device_info device;
|
||||
Kind kind;
|
||||
int group;
|
||||
int64_t k, n, rows, padded;
|
||||
size_t quant, scales, accum, partial, output, gemm_offset;
|
||||
SDI8CudaGemmPlan* gemm = nullptr;
|
||||
@@ -456,7 +501,10 @@ namespace SD_CUDA_NAMESPACE {
|
||||
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;
|
||||
p->group = mat ? op.group_size : 0;
|
||||
if (!mat)
|
||||
sd_ggml_convrot_describe(dst, &p->group);
|
||||
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);
|
||||
@@ -480,7 +528,8 @@ namespace SD_CUDA_NAMESPACE {
|
||||
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_memory;
|
||||
p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
|
||||
p->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
|
||||
: 0);
|
||||
p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
|
||||
*workspace = total;
|
||||
return GGML_STATUS_SUCCESS;
|
||||
@@ -506,6 +555,17 @@ namespace SD_CUDA_NAMESPACE {
|
||||
return error;
|
||||
}
|
||||
const size_t shared = size_t(p.k) * sizeof(float);
|
||||
if (p.group == 1) {
|
||||
quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(src, dst, scales, p.k, p.rows);
|
||||
return cuda_last_error();
|
||||
}
|
||||
if (p.group == 64) {
|
||||
const auto groups = p.k / 64;
|
||||
convrot_h64<false><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
|
||||
reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
|
||||
convrot_h64<true><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
|
||||
return cuda_last_error();
|
||||
}
|
||||
if (shared + 128 <= p.device.shared_memory) {
|
||||
static std::mutex mutex;
|
||||
static std::unordered_map<int, bool> configured;
|
||||
|
||||
@@ -9,6 +9,7 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
void register_rope(ggml_backend_dev_t device);
|
||||
void register_i8(ggml_backend_dev_t device);
|
||||
void register_w4(ggml_backend_dev_t device);
|
||||
void register_convrot(ggml_backend_dev_t device);
|
||||
void register_affine(ggml_backend_dev_t device);
|
||||
void register_modulate(ggml_backend_dev_t device);
|
||||
@@ -25,6 +26,7 @@ namespace SD_CUDA_NAMESPACE {
|
||||
register_affine(device);
|
||||
register_modulate(device);
|
||||
register_i8(device);
|
||||
register_w4(device);
|
||||
register_convrot(device);
|
||||
register_rope(device);
|
||||
register_bf16(device);
|
||||
|
||||
@@ -0,0 +1,873 @@
|
||||
#include <climits>
|
||||
#include "core/ggml_extensions.h"
|
||||
#include "runtime.h"
|
||||
|
||||
namespace SD_CUDA_NAMESPACE {
|
||||
namespace sd_w4_cuda {
|
||||
static __device__ __forceinline__ int shuffle_sum(int value, int lane, int width) {
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
// ROCm LLVM 24 can crash lowering __shfl's wave_shuffle on RDNA2.
|
||||
// These 8/16-lane groups stay within one ds_bpermute half-wave.
|
||||
const int index = (lane & (width - 1)) + (__lane_id() & ~(width - 1));
|
||||
return __builtin_amdgcn_ds_bpermute(index * 4, value);
|
||||
#else
|
||||
return cuda_shuffle(value, lane, width);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Signed byte dots preserve the INT4 grid after the XOR-8 bias correction.
|
||||
static __device__ __forceinline__ int dot4(int a, int b, int c) {
|
||||
#if defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
|
||||
return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
|
||||
#elif defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__))
|
||||
return __builtin_amdgcn_sdot4(a, b, c, false);
|
||||
#elif defined(SD_CUDA_VENDOR_MUSA) || (!defined(SD_CUDA_VENDOR_HIP) && __CUDA_ARCH__ >= 610)
|
||||
return __dp4a(a, b, c);
|
||||
#else
|
||||
const auto* av = reinterpret_cast<const signed char*>(&a);
|
||||
const auto* bv = reinterpret_cast<const signed char*>(&b);
|
||||
for (int i = 0; i < 4; ++i)
|
||||
c += int(av[i]) * int(bv[i]);
|
||||
return c;
|
||||
#endif
|
||||
}
|
||||
static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
|
||||
const int exponent = (v >> 3) & 15, mantissa = v & 7;
|
||||
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;
|
||||
}
|
||||
// A CTA shares decoded weights across 32 activation rows. Each double-buffered
|
||||
// stage covers 64 K values; W4A8 restores the codebook scale per 16-value group.
|
||||
|
||||
static __device__ __forceinline__ uint32_t perm_bytes(uint32_t lo, uint32_t hi, uint32_t sel) {
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
return __builtin_amdgcn_perm(hi, lo, sel);
|
||||
#elif !defined(SD_CUDA_VENDOR_MUSA)
|
||||
return __byte_perm(lo, hi, sel);
|
||||
#else
|
||||
uint32_t r = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const uint32_t b = (sel >> (4 * i)) & 0xF;
|
||||
const uint32_t byte = (b < 4 ? lo : hi) >> (8 * (b & 3)) & 0xFF;
|
||||
r |= byte << (8 * i);
|
||||
}
|
||||
return r;
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(SD_CUDA_VENDOR_HIP)
|
||||
#define PERM_SEL_EVEN 0x06040200u
|
||||
#define PERM_SEL_ODD 0x07050301u
|
||||
#else
|
||||
#define PERM_SEL_EVEN 0x6420u
|
||||
#define PERM_SEL_ODD 0x7531u
|
||||
#endif
|
||||
|
||||
static __device__ __forceinline__ void interleave_act_pair(
|
||||
const int lo,
|
||||
const int hi,
|
||||
int* __restrict__ dst) {
|
||||
dst[0] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_EVEN);
|
||||
dst[1] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_ODD);
|
||||
}
|
||||
|
||||
#define MUL_MAT_W4_BLOCK 16
|
||||
#define MUL_MAT_W4_ROWS (2 * MUL_MAT_W4_BLOCK)
|
||||
#define MUL_MAT_W4_GPS 4
|
||||
#define MUL_MAT_W4_SLOTS (4 * MUL_MAT_W4_GPS)
|
||||
#define MUL_MAT_W4_OUT (MUL_MAT_W4_BLOCK * 4)
|
||||
|
||||
static __global__ void __launch_bounds__(MUL_MAT_W4_BLOCK* MUL_MAT_W4_BLOCK) mul_mat_w4a4_convrot_dp4a(
|
||||
const char* __restrict__ w,
|
||||
const char* __restrict__ x,
|
||||
int* __restrict__ y,
|
||||
const uint32_t n,
|
||||
const uint32_t rows,
|
||||
const uint32_t packed_k) {
|
||||
const uint32_t o_base = blockIdx.x * MUL_MAT_W4_OUT;
|
||||
const uint32_t r0 = blockIdx.y * MUL_MAT_W4_ROWS + threadIdx.y;
|
||||
const uint32_t r1 = r0 + MUL_MAT_W4_BLOCK;
|
||||
const uint32_t lx = threadIdx.x;
|
||||
const uint32_t ly = threadIdx.y;
|
||||
const uint32_t idx = ly * MUL_MAT_W4_BLOCK + lx;
|
||||
const uint32_t dc = idx & (MUL_MAT_W4_OUT - 1);
|
||||
const uint32_t dw = idx >> 6;
|
||||
|
||||
__shared__ int w_dec[2][MUL_MAT_W4_SLOTS * MUL_MAT_W4_OUT];
|
||||
__shared__ int x_tile[2][MUL_MAT_W4_ROWS * MUL_MAT_W4_SLOTS];
|
||||
|
||||
const uint32_t w_words = packed_k / 4;
|
||||
const uint32_t x_words = (packed_k * 2) / 4;
|
||||
const uint32_t groups = packed_k / 8;
|
||||
const int* wi = (const int*)w;
|
||||
const int* xi = (const int*)x;
|
||||
|
||||
const uint32_t fl = lx & 7;
|
||||
const uint32_t fr = (lx >> 3) ? r1 : r0;
|
||||
const uint32_t tr = (lx >> 3) * MUL_MAT_W4_BLOCK + ly;
|
||||
int my_sum = 0;
|
||||
|
||||
{
|
||||
const uint32_t ocol = o_base + dc;
|
||||
int2 wv = make_int2(0, 0);
|
||||
if (ocol < n && dw * 2 < w_words) {
|
||||
wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + dw];
|
||||
}
|
||||
w_dec[0][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[0][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[0][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[0][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
|
||||
|
||||
const uint32_t gi = fl >> 1;
|
||||
const uint32_t xb = gi * 4 + (fl & 1) * 2;
|
||||
const int a = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
|
||||
const int b = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
|
||||
interleave_act_pair(a, b, &x_tile[0][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
|
||||
my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int acc0a = 0, acc0b = 0;
|
||||
int acc1a = 0, acc1b = 0;
|
||||
int acc2a = 0, acc2b = 0;
|
||||
int acc3a = 0, acc3b = 0;
|
||||
int acc4a = 0, acc4b = 0;
|
||||
int acc5a = 0, acc5b = 0;
|
||||
int acc6a = 0, acc6b = 0;
|
||||
int acc7a = 0, acc7b = 0;
|
||||
|
||||
for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4_GPS) {
|
||||
const uint32_t buf = (gs >> 2) & 1;
|
||||
|
||||
if (gs + MUL_MAT_W4_GPS < groups) {
|
||||
const uint32_t nbuf = buf ^ 1;
|
||||
const uint32_t ngs = gs + MUL_MAT_W4_GPS;
|
||||
{
|
||||
const uint32_t ocol = o_base + dc;
|
||||
const uint32_t widx = ngs * 2 + dw * 2;
|
||||
int2 wv = make_int2(0, 0);
|
||||
if (ocol < n && widx < w_words) {
|
||||
wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + ngs + dw];
|
||||
}
|
||||
w_dec[nbuf][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[nbuf][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[nbuf][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
|
||||
w_dec[nbuf][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
|
||||
const uint32_t gi = ngs + (fl >> 1);
|
||||
const uint32_t xb = gi * 4 + (fl & 1) * 2;
|
||||
const int a = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
|
||||
const int b = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
|
||||
interleave_act_pair(a, b, &x_tile[nbuf][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
|
||||
my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t gl = 0; gl < MUL_MAT_W4_GPS; ++gl) {
|
||||
const int u0[4] = {
|
||||
x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 0],
|
||||
x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 1],
|
||||
x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 2],
|
||||
x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 3],
|
||||
};
|
||||
const int u1[4] = {
|
||||
x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 0],
|
||||
x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 1],
|
||||
x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 2],
|
||||
x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 3],
|
||||
};
|
||||
#pragma unroll
|
||||
for (uint32_t jp = 0; jp < 2; ++jp) {
|
||||
const int ue0 = jp == 0 ? u0[0] : u0[2];
|
||||
const int uo0 = jp == 0 ? u0[1] : u0[3];
|
||||
const int ue1 = jp == 0 ? u1[0] : u1[2];
|
||||
const int uo1 = jp == 0 ? u1[1] : u1[3];
|
||||
const uint32_t we = (gl * 4 + jp * 2) * MUL_MAT_W4_OUT;
|
||||
const uint32_t wo = we + MUL_MAT_W4_OUT;
|
||||
|
||||
acc0a = dot4(w_dec[buf][we + lx], ue0, acc0a);
|
||||
acc1a = dot4(w_dec[buf][we + 16 + lx], ue0, acc1a);
|
||||
acc2a = dot4(w_dec[buf][we + 32 + lx], ue0, acc2a);
|
||||
acc3a = dot4(w_dec[buf][we + 48 + lx], ue0, acc3a);
|
||||
acc0b = dot4(w_dec[buf][wo + lx], uo0, acc0b);
|
||||
acc1b = dot4(w_dec[buf][wo + 16 + lx], uo0, acc1b);
|
||||
acc2b = dot4(w_dec[buf][wo + 32 + lx], uo0, acc2b);
|
||||
acc3b = dot4(w_dec[buf][wo + 48 + lx], uo0, acc3b);
|
||||
acc4a = dot4(w_dec[buf][we + lx], ue1, acc4a);
|
||||
acc5a = dot4(w_dec[buf][we + 16 + lx], ue1, acc5a);
|
||||
acc6a = dot4(w_dec[buf][we + 32 + lx], ue1, acc6a);
|
||||
acc7a = dot4(w_dec[buf][we + 48 + lx], ue1, acc7a);
|
||||
acc4b = dot4(w_dec[buf][wo + lx], uo1, acc4b);
|
||||
acc5b = dot4(w_dec[buf][wo + 16 + lx], uo1, acc5b);
|
||||
acc6b = dot4(w_dec[buf][wo + 32 + lx], uo1, acc6b);
|
||||
acc7b = dot4(w_dec[buf][wo + 48 + lx], uo1, acc7b);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
int v = my_sum;
|
||||
v += shuffle_sum(v, (int(threadIdx.x) % 8 + 4) % 8, 8);
|
||||
v += shuffle_sum(v, (int(threadIdx.x) % 8 + 2) % 8, 8);
|
||||
v += shuffle_sum(v, (int(threadIdx.x) % 8 + 1) % 8, 8);
|
||||
const int sum0 = shuffle_sum(v, 0, 16);
|
||||
const int sum1 = shuffle_sum(v, 8, 16);
|
||||
|
||||
const int accv[2][4] = {
|
||||
{acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
|
||||
{acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
|
||||
};
|
||||
const int sumv[2] = {sum0, sum1};
|
||||
#pragma unroll
|
||||
for (uint32_t t = 0; t < 2; ++t) {
|
||||
const uint32_t rt = t == 0 ? r0 : r1;
|
||||
if (rt < rows) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
const uint32_t oj = o_base + j * 16 + lx;
|
||||
if (oj < n) {
|
||||
y[(size_t)rt * n + oj] = accv[t][j] - 8 * sumv[t];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define MUL_MAT_W4A8_BLOCK 16
|
||||
#define MUL_MAT_W4A8_ROWS (2 * MUL_MAT_W4A8_BLOCK)
|
||||
#define MUL_MAT_W4A8_GPS 4
|
||||
#define MUL_MAT_W4A8_SLOTS (4 * MUL_MAT_W4A8_GPS)
|
||||
#define MUL_MAT_W4A8_OUT (MUL_MAT_W4A8_BLOCK * 4)
|
||||
|
||||
#if defined(SD_CUDA_VENDOR_HIP) && defined(__HIP_DEVICE_COMPILE__)
|
||||
// __has_builtin checks the active device target's features, not just whether
|
||||
// the compiler recognizes the name. Prefer the non-replicated fragment layout.
|
||||
#if __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12)
|
||||
#define SD_W4_WMMA_GFX12
|
||||
#elif __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64)
|
||||
#define SD_W4_WMMA_GFX11
|
||||
#endif
|
||||
#endif
|
||||
|
||||
// Keep each intrinsic paired with its fragment layout, including on newer GPUs.
|
||||
#if defined(SD_W4_WMMA_GFX11)
|
||||
typedef int v4si32_w4a8 __attribute__((ext_vector_type(4)));
|
||||
|
||||
static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma(
|
||||
const char* __restrict__ w,
|
||||
const char* __restrict__ x,
|
||||
const float* __restrict__ codebook,
|
||||
const float* __restrict__ s_channel,
|
||||
const uint8_t* __restrict__ s_rel,
|
||||
const float* __restrict__ act_scales,
|
||||
const float* __restrict__ bias,
|
||||
float* __restrict__ y,
|
||||
const uint32_t n,
|
||||
const uint32_t rows,
|
||||
const uint32_t k,
|
||||
const uint32_t r_base) {
|
||||
const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
|
||||
const uint32_t idx = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
|
||||
const uint32_t wv = idx >> 6;
|
||||
const uint32_t lane = idx & 63;
|
||||
const uint32_t j = lane & 15;
|
||||
const uint32_t q = lane >> 4;
|
||||
const uint32_t ocol = o_base + wv * 16 + j;
|
||||
const uint32_t ra = r_base + j;
|
||||
|
||||
__shared__ __align__(4) int8_t cb8[16];
|
||||
__shared__ float s_cb;
|
||||
__shared__ int lut[256];
|
||||
__shared__ float s_lut[256];
|
||||
|
||||
const uint32_t w_words = k / 8;
|
||||
const uint32_t x_words = k / 4;
|
||||
const uint32_t groups = k / 16;
|
||||
const int* wi = (const int*)w;
|
||||
const int* xi = (const int*)x;
|
||||
|
||||
if (idx == 0) {
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
amax = fmaxf(amax, fabsf(codebook[i]));
|
||||
}
|
||||
s_cb = amax / 127.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
|
||||
cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
s_lut[b] = fp8_e4m3(uint8_t(b));
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const bool o_valid = ocol < n;
|
||||
const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
|
||||
const float scb = s_cb;
|
||||
|
||||
float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
||||
const v4si32_w4a8 zero = {};
|
||||
|
||||
int w0c = 0, w1c = 0, xc0 = 0, xc1 = 0, xc2 = 0, xc3 = 0, sc = 0;
|
||||
int w0n = 0, w1n = 0, xn0 = 0, xn1 = 0, xn2 = 0, xn3 = 0, sn = 0;
|
||||
{
|
||||
const int* xr = xi + (size_t)ra * x_words;
|
||||
xc0 = ra < rows ? xr[0] : 0;
|
||||
xc1 = ra < rows ? xr[1] : 0;
|
||||
xc2 = ra < rows ? xr[2] : 0;
|
||||
xc3 = ra < rows ? xr[3] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words;
|
||||
w0c = wp[0];
|
||||
w1c = wp[1];
|
||||
sc = s_rel[(size_t)ocol * groups];
|
||||
}
|
||||
}
|
||||
if (groups > 1) {
|
||||
const int* xr = xi + (size_t)ra * x_words + 4;
|
||||
xn0 = ra < rows ? xr[0] : 0;
|
||||
xn1 = ra < rows ? xr[1] : 0;
|
||||
xn2 = ra < rows ? xr[2] : 0;
|
||||
xn3 = ra < rows ? xr[3] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words + 2;
|
||||
w0n = wp[0];
|
||||
w1n = wp[1];
|
||||
sn = s_rel[(size_t)ocol * groups + 1];
|
||||
}
|
||||
}
|
||||
|
||||
for (uint32_t g = 0; g < groups; ++g) {
|
||||
v4si32_w4a8 bfrag = {};
|
||||
float s = 0.0f;
|
||||
if (o_valid) {
|
||||
bfrag = {
|
||||
(lut[w0c & 0xFF] & 0xFFFF) | ((lut[(w0c >> 8) & 0xFF] & 0xFFFF) << 16),
|
||||
(lut[(w0c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w0c >> 24) & 0xFF] & 0xFFFF) << 16),
|
||||
(lut[w1c & 0xFF] & 0xFFFF) | ((lut[(w1c >> 8) & 0xFF] & 0xFFFF) << 16),
|
||||
(lut[(w1c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w1c >> 24) & 0xFF] & 0xFFFF) << 16),
|
||||
};
|
||||
|
||||
s = scb * (s_chan * s_lut[sc & 0xFF]);
|
||||
}
|
||||
const v4si32_w4a8 afrag = {xc0, xc1, xc2, xc3};
|
||||
const v4si32_w4a8 dfrag =
|
||||
__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64(true, afrag, true, bfrag, zero, false);
|
||||
|
||||
w0c = w0n;
|
||||
w1c = w1n;
|
||||
sc = sn;
|
||||
xc0 = xn0;
|
||||
xc1 = xn1;
|
||||
xc2 = xn2;
|
||||
xc3 = xn3;
|
||||
if (g + 2 < groups) {
|
||||
const int* xr = xi + (size_t)ra * x_words + (g + 2) * 4;
|
||||
xn0 = ra < rows ? xr[0] : 0;
|
||||
xn1 = ra < rows ? xr[1] : 0;
|
||||
xn2 = ra < rows ? xr[2] : 0;
|
||||
xn3 = ra < rows ? xr[3] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
|
||||
w0n = wp[0];
|
||||
w1n = wp[1];
|
||||
sn = s_rel[(size_t)ocol * groups + g + 2];
|
||||
}
|
||||
}
|
||||
|
||||
acc[0] = fmaf((float)dfrag.x, s, acc[0]);
|
||||
acc[1] = fmaf((float)dfrag.y, s, acc[1]);
|
||||
acc[2] = fmaf((float)dfrag.z, s, acc[2]);
|
||||
acc[3] = fmaf((float)dfrag.w, s, acc[3]);
|
||||
}
|
||||
|
||||
if (o_valid) {
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < 4; ++d) {
|
||||
const uint32_t r = r_base + 4 * d + q;
|
||||
if (r < rows) {
|
||||
float value = acc[d] * act_scales[r];
|
||||
if (bias != nullptr) {
|
||||
value += bias[ocol];
|
||||
}
|
||||
y[(size_t)r * n + ocol] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
#if defined(SD_W4_WMMA_GFX12)
|
||||
typedef int v4si32_w4a8_g12 __attribute__((ext_vector_type(4)));
|
||||
|
||||
static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma_gfx12(
|
||||
const char* __restrict__ w,
|
||||
const char* __restrict__ x,
|
||||
const float* __restrict__ codebook,
|
||||
const float* __restrict__ s_channel,
|
||||
const uint8_t* __restrict__ s_rel,
|
||||
const float* __restrict__ act_scales,
|
||||
const float* __restrict__ bias,
|
||||
float* __restrict__ y,
|
||||
const uint32_t n,
|
||||
const uint32_t rows,
|
||||
const uint32_t k,
|
||||
const uint32_t r_base) {
|
||||
const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
|
||||
const uint32_t idx = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
|
||||
const uint32_t wv = idx >> 6;
|
||||
const uint32_t lane = idx & 63;
|
||||
const uint32_t j = lane & 15;
|
||||
const uint32_t h = lane >> 4;
|
||||
const uint32_t kh = ((h & 1) << 1) | (h >> 1);
|
||||
const uint32_t ocol = o_base + wv * 16 + j;
|
||||
const uint32_t ra = r_base + j;
|
||||
|
||||
__shared__ __align__(4) int8_t cb8[16];
|
||||
__shared__ float s_cb;
|
||||
__shared__ int lut[256];
|
||||
__shared__ float s_lut[256];
|
||||
|
||||
const uint32_t w_words = k / 8;
|
||||
const uint32_t x_words = k / 4;
|
||||
const uint32_t groups = k / 16;
|
||||
const int* wi = (const int*)w;
|
||||
const int* xi = (const int*)x;
|
||||
|
||||
if (idx == 0) {
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
amax = fmaxf(amax, fabsf(codebook[i]));
|
||||
}
|
||||
s_cb = amax / 127.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
|
||||
cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
s_lut[b] = fp8_e4m3(uint8_t(b));
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const bool o_valid = ocol < n;
|
||||
const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
|
||||
const float scb = s_cb;
|
||||
|
||||
float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
|
||||
const v4si32_w4a8_g12 zero = {};
|
||||
|
||||
int w0c = 0, w1c = 0, xc = 0, sc = 0;
|
||||
int w0n = 0, w1n = 0, xn = 0, sn = 0;
|
||||
{
|
||||
const int* xr = xi + (size_t)ra * x_words;
|
||||
xc = ra < rows ? xr[kh] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words;
|
||||
w0c = wp[0];
|
||||
w1c = wp[1];
|
||||
sc = s_rel[(size_t)ocol * groups];
|
||||
}
|
||||
}
|
||||
if (groups > 1) {
|
||||
const int* xr = xi + (size_t)ra * x_words;
|
||||
xn = ra < rows ? xr[4 + kh] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words + 2;
|
||||
w0n = wp[0];
|
||||
w1n = wp[1];
|
||||
sn = s_rel[(size_t)ocol * groups + 1];
|
||||
}
|
||||
}
|
||||
|
||||
for (uint32_t g = 0; g < groups; ++g) {
|
||||
int bfrag = 0;
|
||||
float s = 0.0f;
|
||||
if (o_valid) {
|
||||
const int bwd = (kh & 2) ? w1c : w0c;
|
||||
const int bsh = (kh & 1) * 16;
|
||||
bfrag = (lut[(bwd >> bsh) & 0xFF] & 0xFFFF) |
|
||||
((lut[(bwd >> (bsh + 8)) & 0xFF] & 0xFFFF) << 16);
|
||||
|
||||
s = scb * (s_chan * s_lut[sc & 0xFF]);
|
||||
}
|
||||
const v4si32_w4a8_g12 dfrag =
|
||||
__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12(true, xc, true, bfrag, zero, false);
|
||||
|
||||
w0c = w0n;
|
||||
w1c = w1n;
|
||||
sc = sn;
|
||||
xc = xn;
|
||||
if (g + 2 < groups) {
|
||||
const int* xr = xi + (size_t)ra * x_words;
|
||||
xn = ra < rows ? xr[(g + 2) * 4 + kh] : 0;
|
||||
if (o_valid) {
|
||||
const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
|
||||
w0n = wp[0];
|
||||
w1n = wp[1];
|
||||
sn = s_rel[(size_t)ocol * groups + g + 2];
|
||||
}
|
||||
}
|
||||
|
||||
acc[0] = fmaf((float)dfrag.x, s, acc[0]);
|
||||
acc[1] = fmaf((float)dfrag.y, s, acc[1]);
|
||||
acc[2] = fmaf((float)dfrag.z, s, acc[2]);
|
||||
acc[3] = fmaf((float)dfrag.w, s, acc[3]);
|
||||
}
|
||||
|
||||
if (o_valid) {
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < 4; ++d) {
|
||||
const uint32_t r = r_base + 4 * kh + d;
|
||||
if (r < rows) {
|
||||
float value = acc[d] * act_scales[r];
|
||||
if (bias != nullptr) {
|
||||
value += bias[ocol];
|
||||
}
|
||||
y[(size_t)r * n + ocol] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
static __global__ void __launch_bounds__(MUL_MAT_W4A8_BLOCK* MUL_MAT_W4A8_BLOCK) mul_mat_w4a8_convrot(
|
||||
const char* __restrict__ w,
|
||||
const char* __restrict__ x,
|
||||
const float* __restrict__ codebook,
|
||||
const float* __restrict__ s_channel,
|
||||
const uint8_t* __restrict__ s_rel,
|
||||
const float* __restrict__ act_scales,
|
||||
const float* __restrict__ bias,
|
||||
float* __restrict__ y,
|
||||
const uint32_t n,
|
||||
const uint32_t rows,
|
||||
const uint32_t k) {
|
||||
#if defined(SD_W4_WMMA_GFX11) || defined(SD_W4_WMMA_GFX12)
|
||||
// Retain the 32-row launch grid shared with DP4A. Each WMMA wave
|
||||
// computes a 16-row tile; the barrier protects the reused LUTs.
|
||||
for (uint32_t tile = 0; tile < 2; ++tile) {
|
||||
#if defined(SD_W4_WMMA_GFX12)
|
||||
mul_mat_w4a8_convrot_wmma_gfx12(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
|
||||
blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
|
||||
#else
|
||||
mul_mat_w4a8_convrot_wmma(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
|
||||
blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
|
||||
#endif
|
||||
__syncthreads();
|
||||
}
|
||||
#else
|
||||
const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
|
||||
const uint32_t r0 = blockIdx.y * MUL_MAT_W4A8_ROWS + threadIdx.y;
|
||||
const uint32_t r1 = r0 + MUL_MAT_W4A8_BLOCK;
|
||||
const uint32_t lx = threadIdx.x;
|
||||
const uint32_t ly = threadIdx.y;
|
||||
const uint32_t idx = ly * MUL_MAT_W4A8_BLOCK + lx;
|
||||
const uint32_t dc = idx & (MUL_MAT_W4A8_OUT - 1);
|
||||
const uint32_t dw = idx >> 6;
|
||||
|
||||
__shared__ int w_dec[2][MUL_MAT_W4A8_SLOTS * MUL_MAT_W4A8_OUT];
|
||||
__shared__ int x_tile[2][MUL_MAT_W4A8_ROWS * MUL_MAT_W4A8_SLOTS];
|
||||
__shared__ float s_dec[2][MUL_MAT_W4A8_GPS * MUL_MAT_W4A8_OUT];
|
||||
__shared__ __align__(4) int8_t cb8[16];
|
||||
__shared__ float s_cb;
|
||||
__shared__ int lut[256];
|
||||
|
||||
const uint32_t w_words = k / 8;
|
||||
const uint32_t x_words = k / 4;
|
||||
const uint32_t groups = k / 16;
|
||||
const int* wi = (const int*)w;
|
||||
const int* xi = (const int*)x;
|
||||
|
||||
if (ly == 0 && lx == 0) {
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
amax = fmaxf(amax, fabsf(codebook[i]));
|
||||
}
|
||||
|
||||
s_cb = amax / 127.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
const int v = s_cb == 0.f ? 0 : __float2int_rn(codebook[i] / s_cb);
|
||||
cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
lut[idx] = (cb8[idx & 0xF] & 0xFF) | ((cb8[idx >> 4] & 0xFF) << 8);
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 2; ++j) {
|
||||
const uint32_t pw = dw * 2 + j;
|
||||
const uint32_t ocol = o_base + dc;
|
||||
const int wv = (ocol < n && pw < w_words) ? wi[(size_t)ocol * w_words + pw] : 0;
|
||||
w_dec[0][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
|
||||
w_dec[0][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
|
||||
}
|
||||
{
|
||||
const uint32_t ocol = o_base + dc;
|
||||
float s = 0.0f;
|
||||
if (ocol < n && dw < groups) {
|
||||
s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + dw]);
|
||||
}
|
||||
s_dec[0][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
|
||||
}
|
||||
if (lx < 16) {
|
||||
const uint32_t aidx = lx;
|
||||
x_tile[0][ly * MUL_MAT_W4A8_SLOTS + lx] =
|
||||
(r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
|
||||
x_tile[0][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
|
||||
(r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
float acc0 = 0.0f;
|
||||
float acc1 = 0.0f;
|
||||
float acc2 = 0.0f;
|
||||
float acc3 = 0.0f;
|
||||
float acc4 = 0.0f;
|
||||
float acc5 = 0.0f;
|
||||
float acc6 = 0.0f;
|
||||
float acc7 = 0.0f;
|
||||
|
||||
for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4A8_GPS) {
|
||||
const uint32_t buf = (gs >> 2) & 1;
|
||||
|
||||
if (gs + MUL_MAT_W4A8_GPS < groups) {
|
||||
const uint32_t nbuf = buf ^ 1;
|
||||
const uint32_t ngs = gs + MUL_MAT_W4A8_GPS;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 2; ++j) {
|
||||
const uint32_t pw = dw * 2 + j;
|
||||
const uint32_t ocol = o_base + dc;
|
||||
const uint32_t widx = ngs * 2 + pw;
|
||||
const int wv = (ocol < n && widx < w_words) ? wi[(size_t)ocol * w_words + widx] : 0;
|
||||
w_dec[nbuf][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
|
||||
w_dec[nbuf][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
|
||||
}
|
||||
{
|
||||
const uint32_t ocol = o_base + dc;
|
||||
float s = 0.0f;
|
||||
if (ocol < n && ngs + dw < groups) {
|
||||
s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + ngs + dw]);
|
||||
}
|
||||
s_dec[nbuf][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
|
||||
}
|
||||
if (lx < 16) {
|
||||
const uint32_t aidx = ngs * 4 + lx;
|
||||
x_tile[nbuf][ly * MUL_MAT_W4A8_SLOTS + lx] =
|
||||
(r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
|
||||
x_tile[nbuf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
|
||||
(r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t gl = 0; gl < MUL_MAT_W4A8_GPS; ++gl) {
|
||||
const int u0[4] = {
|
||||
x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
|
||||
x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
|
||||
x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
|
||||
x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
|
||||
};
|
||||
const int u1[4] = {
|
||||
x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
|
||||
x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
|
||||
x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
|
||||
x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
|
||||
};
|
||||
int accg0 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u0[0], 0);
|
||||
accg0 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u0[1], accg0);
|
||||
accg0 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u0[2], accg0);
|
||||
accg0 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u0[3], accg0);
|
||||
int accg1 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u0[0], 0);
|
||||
accg1 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u0[1], accg1);
|
||||
accg1 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u0[2], accg1);
|
||||
accg1 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u0[3], accg1);
|
||||
int accg2 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u0[0], 0);
|
||||
accg2 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u0[1], accg2);
|
||||
accg2 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u0[2], accg2);
|
||||
accg2 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u0[3], accg2);
|
||||
int accg3 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u0[0], 0);
|
||||
accg3 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u0[1], accg3);
|
||||
accg3 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u0[2], accg3);
|
||||
accg3 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u0[3], accg3);
|
||||
int accg4 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u1[0], 0);
|
||||
accg4 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u1[1], accg4);
|
||||
accg4 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u1[2], accg4);
|
||||
accg4 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u1[3], accg4);
|
||||
int accg5 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u1[0], 0);
|
||||
accg5 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u1[1], accg5);
|
||||
accg5 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u1[2], accg5);
|
||||
accg5 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u1[3], accg5);
|
||||
int accg6 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u1[0], 0);
|
||||
accg6 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u1[1], accg6);
|
||||
accg6 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u1[2], accg6);
|
||||
accg6 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u1[3], accg6);
|
||||
int accg7 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u1[0], 0);
|
||||
accg7 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u1[1], accg7);
|
||||
accg7 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u1[2], accg7);
|
||||
accg7 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u1[3], accg7);
|
||||
acc0 = fmaf((float)accg0, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc0);
|
||||
acc1 = fmaf((float)accg1, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc1);
|
||||
acc2 = fmaf((float)accg2, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc2);
|
||||
acc3 = fmaf((float)accg3, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc3);
|
||||
acc4 = fmaf((float)accg4, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc4);
|
||||
acc5 = fmaf((float)accg5, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc5);
|
||||
acc6 = fmaf((float)accg6, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc6);
|
||||
acc7 = fmaf((float)accg7, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc7);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const float accv[8] = {acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7};
|
||||
#pragma unroll
|
||||
for (uint32_t t = 0; t < 2; ++t) {
|
||||
const uint32_t rt = t == 0 ? r0 : r1;
|
||||
if (rt < rows) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
const uint32_t oj = o_base + j * 16 + lx;
|
||||
if (oj < n) {
|
||||
float value = accv[t * 4 + j] * act_scales[rt];
|
||||
if (bias != nullptr) {
|
||||
value += bias[oj];
|
||||
}
|
||||
y[(size_t)rt * n + oj] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
static __global__ void dequantize(const int32_t* accum, const float* xs, const float* ws, const float* bias, float* dst, int64_t n, int64_t total) {
|
||||
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
if (i < total) {
|
||||
float value = float(accum[i]) * xs[i / n] * ws[i % n];
|
||||
dst[i] = value + (bias ? bias[i % n] : 0.f);
|
||||
}
|
||||
}
|
||||
static bool aligned(const ggml_tensor* t, size_t alignment) {
|
||||
if (!t)
|
||||
return true;
|
||||
size_t offset = 0;
|
||||
while (!t->data && t->view_src) {
|
||||
offset += t->view_offs;
|
||||
t = t->view_src;
|
||||
}
|
||||
return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
|
||||
}
|
||||
static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
|
||||
SDW4Operation op;
|
||||
if (!sd_ggml_w4_describe(node, &op) || cuda_device_info(device).device < 0)
|
||||
return false;
|
||||
const int64_t k = op.input->ne[0], n = node->ne[0], rows = ggml_nrows(node);
|
||||
if (k % 64 || k > INT_MAX || n > INT_MAX - 63 || rows > 65535 * 32 ||
|
||||
ggml_nelements(node) > int64_t(INT_MAX) * 256 || !aligned(op.weight, 8) || !aligned(node, 4))
|
||||
return false;
|
||||
for (auto* t : node->src)
|
||||
if (!aligned(t, 4))
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
|
||||
SDW4Operation op;
|
||||
if (!sd_ggml_w4_describe(nodes[0], &op))
|
||||
return GGML_STATUS_FAILED;
|
||||
const auto* dst = nodes[0];
|
||||
bool may_alias = dst->data || dst->view_src;
|
||||
if (dst->data && dst->buffer && !dst->view_src) {
|
||||
may_alias = false;
|
||||
for (const auto* input : dst->src) {
|
||||
if (!input)
|
||||
continue;
|
||||
size_t offset = 0;
|
||||
const size_t size = ggml_nbytes(input);
|
||||
while (!input->data && input->view_src) {
|
||||
offset += input->view_offs;
|
||||
input = input->view_src;
|
||||
}
|
||||
if (!input->data)
|
||||
continue;
|
||||
const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(input->data) + offset;
|
||||
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));
|
||||
return GGML_STATUS_SUCCESS;
|
||||
}
|
||||
static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
|
||||
if (!src)
|
||||
return false;
|
||||
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*) {
|
||||
auto* dst = nodes[0];
|
||||
SDW4Operation op;
|
||||
ggml_ext_gpu_context context;
|
||||
if (!sd_ggml_w4_describe(dst, &op) || !cuda_context(e->backend, &context))
|
||||
return GGML_STATUS_FAILED;
|
||||
const auto stream = static_cast<CUDAStream>(context.stream);
|
||||
const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0];
|
||||
auto* scratch = e->workspace ? static_cast<float*>(ggml_backend_buffer_get_base(e->workspace)) : nullptr;
|
||||
bool alias = false;
|
||||
for (auto* input : dst->src)
|
||||
alias |= overlaps(dst, input);
|
||||
auto* output = alias ? scratch : static_cast<float*>(dst->data);
|
||||
const auto* w = static_cast<const char*>(op.weight->data);
|
||||
const auto* x = static_cast<const char*>(op.input->data);
|
||||
const auto* xs = reinterpret_cast<const float*>(x + k * ((rows + 3) / 4 * 4));
|
||||
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) {
|
||||
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)
|
||||
return GGML_STATUS_FAILED;
|
||||
dequantize<<<unsigned((n * rows + 255) / 256), 256, 0, stream>>>(accum, xs, ws, bias, output, n, n * rows);
|
||||
} else {
|
||||
mul_mat_w4a8_convrot<<<grid, block, 0, stream>>>(w, x, ws,
|
||||
static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data),
|
||||
xs, bias, output, uint32_t(n), uint32_t(rows), uint32_t(k));
|
||||
}
|
||||
auto error = cuda_last_error();
|
||||
if (error == cuda_success && alias)
|
||||
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
|
||||
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
|
||||
}
|
||||
}
|
||||
void register_w4(ggml_backend_dev_t device) {
|
||||
ggml_ext_kernel_desc d = {};
|
||||
d.abi_version = GGML_EXT_ABI_VERSION;
|
||||
d.name = "sd.linear_w4_convrot." SD_CUDA_NAME ".v1";
|
||||
d.op = GGML_OP_CUSTOM;
|
||||
d.custom_op = sd_ggml_w4_op();
|
||||
d.priority = 100;
|
||||
d.supports = sd_w4_cuda::supports;
|
||||
d.prepare = sd_w4_cuda::prepare;
|
||||
d.execute = sd_w4_cuda::execute;
|
||||
d.destroy = [](void*, void*) {};
|
||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||
}
|
||||
}
|
||||
@@ -301,6 +301,101 @@ ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int gro
|
||||
return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
|
||||
}
|
||||
|
||||
struct W4Attributes {
|
||||
int64_t shape[3];
|
||||
SDW4Kind kind;
|
||||
int32_t has_bias;
|
||||
};
|
||||
|
||||
static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
|
||||
if (!attributes || count < 3)
|
||||
return false;
|
||||
W4Attributes attr;
|
||||
std::memcpy(&attr, attributes, sizeof(attr));
|
||||
const bool codebook = attr.kind == SDW4Kind::W4A8;
|
||||
if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) ||
|
||||
count != (codebook ? 5 : 3) + attr.has_bias)
|
||||
return false;
|
||||
for (int i = 0; i < count; ++i)
|
||||
if (!inputs[i] || !ggml_is_contiguous(inputs[i]))
|
||||
return false;
|
||||
const auto* w = inputs[0];
|
||||
const auto* x = inputs[1];
|
||||
if (w->type != GGML_TYPE_I8 || w->ne[0] <= 0 || w->ne[0] > i8_size_limit / 2 ||
|
||||
w->ne[1] <= 0 || w->ne[2] != 1 || w->ne[3] != 1)
|
||||
return false;
|
||||
const int64_t k = 2 * w->ne[0], n = w->ne[1];
|
||||
int64_t rows = 1;
|
||||
for (auto dim : attr.shape) {
|
||||
if (dim <= 0 || dim > i8_size_limit / rows)
|
||||
return false;
|
||||
rows *= dim;
|
||||
}
|
||||
int64_t height;
|
||||
if (!packed_height(k, rows, &height) || n > i8_size_limit / 4 / rows ||
|
||||
x->type != GGML_TYPE_I8 || x->ne[0] != k || x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1 ||
|
||||
!scale_matches(inputs[2], codebook ? 16 : n) ||
|
||||
(attr.has_bias && !scale_matches(inputs[count - 1], n)))
|
||||
return false;
|
||||
if (codebook && (k % 16 || !scale_matches(inputs[3], n) || inputs[4]->type != GGML_TYPE_F8_E4M3 ||
|
||||
inputs[4]->ne[0] != k / 16 || inputs[4]->ne[1] != n || inputs[4]->ne[2] != 1 || inputs[4]->ne[3] != 1))
|
||||
return false;
|
||||
out->type = GGML_TYPE_F32;
|
||||
out->ne[0] = n;
|
||||
std::memcpy(out->ne + 1, attr.shape, sizeof(attr.shape));
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_ext_op_t sd_ggml_w4_op() {
|
||||
static const auto op = [] {
|
||||
const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer};
|
||||
auto result = ggml_ext_register_op(&desc);
|
||||
GGML_ASSERT(result);
|
||||
return result;
|
||||
}();
|
||||
return op;
|
||||
}
|
||||
|
||||
bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
|
||||
if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node))
|
||||
return false;
|
||||
W4Attributes attr;
|
||||
std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
|
||||
const bool codebook = attr.kind == SDW4Kind::W4A8;
|
||||
const int count = (codebook ? 5 : 3) + (attr.has_bias != 0);
|
||||
ggml_ext_tensor_desc out;
|
||||
if (!w4_infer(node->src, count, &attr, &out) || node->type != out.type ||
|
||||
std::memcmp(node->ne, out.ne, sizeof(out.ne)))
|
||||
return false;
|
||||
*op = {node->src[0], node->src[1], node->src[2], codebook ? node->src[3] : nullptr,
|
||||
codebook ? node->src[4] : nullptr, attr.has_bias ? node->src[count - 1] : nullptr, attr.kind};
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) {
|
||||
if (!w || !x || group < 0 || (group && !valid_group(group)))
|
||||
return nullptr;
|
||||
const int quant_group = group ? group : 1;
|
||||
if (x->type == GGML_TYPE_F32)
|
||||
x = ggml_quantize_i8_convrot(ctx, x, quant_group);
|
||||
int packed_group;
|
||||
if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group)
|
||||
return nullptr;
|
||||
const auto* logical = x->src[0];
|
||||
W4Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0};
|
||||
ggml_tensor* inputs[6] = {w, x, ws};
|
||||
int count = 3;
|
||||
if (kind == SDW4Kind::W4A8) {
|
||||
inputs[count++] = sc;
|
||||
inputs[count++] = sr;
|
||||
} else if (sc || sr) {
|
||||
return nullptr;
|
||||
}
|
||||
if (bias)
|
||||
inputs[count++] = bias;
|
||||
return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &attr);
|
||||
}
|
||||
|
||||
static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
|
||||
if (count != 3 || !attributes)
|
||||
return false;
|
||||
|
||||
@@ -39,6 +39,25 @@ bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
|
||||
ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr);
|
||||
ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
|
||||
|
||||
enum class SDW4Kind : int32_t { W4A4 = 1,
|
||||
W4A8 = 2 };
|
||||
|
||||
struct SDW4Operation {
|
||||
const ggml_tensor* weight;
|
||||
const ggml_tensor* input;
|
||||
const ggml_tensor* weight_scale;
|
||||
const ggml_tensor* s_channel;
|
||||
const ggml_tensor* s_rel;
|
||||
const ggml_tensor* bias;
|
||||
SDW4Kind kind;
|
||||
};
|
||||
|
||||
ggml_ext_op_t sd_ggml_w4_op();
|
||||
bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation);
|
||||
// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
|
||||
// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
|
||||
ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size);
|
||||
|
||||
enum ggml_sage_attn_mode {
|
||||
GGML_SAGE_ATTN_AUTO = 0,
|
||||
GGML_SAGE_ATTN_2,
|
||||
|
||||
@@ -53,6 +53,9 @@ sd_vulkan_shader(rms_modulate rms_modulate ON -O)
|
||||
sd_vulkan_shader(rms_rope_apply rms_rope_apply ON)
|
||||
sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16)
|
||||
sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64)
|
||||
sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_subgroup quantize_i8_convrot_subgroup ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_register quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSINGLE_PASS=1)
|
||||
sd_vulkan_shader(swiglu_quantize_i8_convrot quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1)
|
||||
|
||||
@@ -11,7 +11,10 @@
|
||||
#include "mul_mat_i8_tensorwise.h"
|
||||
#include "mul_mat_i8_tensorwise_cm1.h"
|
||||
#include "mul_mat_i8_tensorwise_cm2.h"
|
||||
#include "mul_mat_w4a4_convrot.h"
|
||||
#include "mul_mat_w4a8_convrot.h"
|
||||
#include "quantize_i8_convrot.h"
|
||||
#include "quantize_i8_convrot_h64.h"
|
||||
#include "quantize_i8_convrot_register.h"
|
||||
#include "quantize_i8_convrot_subgroup.h"
|
||||
#include "rms_modulate_convrot.h"
|
||||
@@ -111,6 +114,23 @@ namespace sd_vulkan_i8 {
|
||||
(cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
|
||||
d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
|
||||
}
|
||||
static bool w4_supports(ggml_backend_dev_t dev, const ggml_tensor* dst, void*) {
|
||||
Device d;
|
||||
SDW4Operation op;
|
||||
if (!device_info(dev, d) || !sd_ggml_w4_describe(dst, &op) || !bounded(d, dst) ||
|
||||
sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot) <= 4 || sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot) <= 4 ||
|
||||
d.compute.shared_memory < (op.kind == SDW4Kind::W4A8 ? 16452u : 13312u) || d.info.max_push_constants < 16)
|
||||
return false;
|
||||
for (auto* t : dst->src)
|
||||
if (t && (!bounded(d, t) || ggml_nbytes(t) > UINT32_MAX))
|
||||
return false;
|
||||
const uint64_t k = op.input->ne[0], n = dst->ne[0], rows = ggml_nrows(dst);
|
||||
return k % 64 == 0 && k <= UINT32_MAX && n <= UINT32_MAX && rows <= UINT32_MAX &&
|
||||
ggml_nelements(dst) <= UINT32_MAX &&
|
||||
(k * ((rows + 3) / 4 * 4)) % d.info.buffer_alignment == 0 &&
|
||||
(n + 63) / 64 <= d.info.max_group_count[0] && (rows + 31) / 32 <= d.info.max_group_count[1] &&
|
||||
d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
|
||||
}
|
||||
static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
|
||||
if (i + 2 >= count_nodes(graph) || !device.subgroup) {
|
||||
return false;
|
||||
@@ -534,9 +554,12 @@ namespace sd_vulkan_i8 {
|
||||
status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
|
||||
} else {
|
||||
ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
|
||||
auto pipeline = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
|
||||
bool subgroup = d.subgroup && k >= 1024;
|
||||
uint32_t width = k;
|
||||
int group;
|
||||
if (!sd_ggml_convrot_describe(output, &group))
|
||||
return GGML_STATUS_FAILED;
|
||||
auto pipeline = group == 64 ? SD_VK_PIPELINE(quantize_i8_convrot_h64, 3, QuantParameters) : SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
|
||||
bool subgroup = group == 256 && d.subgroup && k >= 1024;
|
||||
uint32_t width = k;
|
||||
if (subgroup) {
|
||||
pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
|
||||
if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
|
||||
@@ -653,6 +676,42 @@ namespace sd_vulkan_i8 {
|
||||
}
|
||||
return finish(e, d, output, dst, alias, status);
|
||||
}
|
||||
static ggml_status multiply_w4(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* output) {
|
||||
SDW4Operation op;
|
||||
if (!sd_ggml_w4_describe(output, &op))
|
||||
return GGML_STATUS_FAILED;
|
||||
const uint32_t k = uint32_t(op.input->ne[0]), n = uint32_t(output->ne[0]), rows = uint32_t(ggml_nrows(output));
|
||||
const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
|
||||
const ggml_ext_vulkan_binding xs = {op.input, nullptr, scale_offset, size_t(rows) * 4};
|
||||
std::vector<ggml_ext_vulkan_binding> bindings = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}};
|
||||
if (op.kind == SDW4Kind::W4A8) {
|
||||
bindings.push_back({op.s_channel});
|
||||
bindings.push_back({op.s_rel});
|
||||
}
|
||||
bindings.push_back(xs);
|
||||
bindings.push_back({op.bias ? op.bias : op.weight_scale});
|
||||
ggml_ext_vulkan_binding dst;
|
||||
bool alias;
|
||||
if (!output_binding(e, d, output, bindings, dst, alias))
|
||||
return GGML_STATUS_FAILED;
|
||||
bindings.push_back(dst);
|
||||
struct Parameters {
|
||||
uint32_t n, rows, k, bias;
|
||||
};
|
||||
const Parameters pc = {n, rows, k, op.bias ? 1u : 0u};
|
||||
auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
|
||||
const uint32_t groups[] = {(n + 63) / 64, (rows + 31) / 32, 1};
|
||||
auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
|
||||
return finish(e, d, output, dst, alias, status);
|
||||
}
|
||||
static bool quant64_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
|
||||
Device d;
|
||||
int group;
|
||||
return device_info(dev, d) && sd_ggml_convrot_describe(n, &group) && group == 64 &&
|
||||
bounded(d, n) && bounded(d, n->src[0]) && n->ne[0] <= UINT32_MAX &&
|
||||
ggml_nrows(n->src[0]) <= UINT32_MAX && d.info.max_group_invocations >= 16 && d.info.max_group_size[0] >= 16 &&
|
||||
(uint64_t(n->ne[0]) * ((ggml_nrows(n->src[0]) + 3) / 4 * 4)) % d.info.buffer_alignment == 0;
|
||||
}
|
||||
static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
|
||||
Device d;
|
||||
return device_info(dev, d) && can_quantize(d, n);
|
||||
@@ -747,6 +806,8 @@ namespace sd_vulkan_i8 {
|
||||
Device d;
|
||||
if (!device_info(ggml_backend_get_device(e->backend), d))
|
||||
return GGML_STATUS_FAILED;
|
||||
if (ggml_ext_op_get(nodes[0]) == sd_ggml_w4_op())
|
||||
return multiply_w4(e, d, nodes[0]);
|
||||
return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
|
||||
}
|
||||
#undef SD_VK_PIPELINE
|
||||
@@ -768,6 +829,13 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
|
||||
d.custom_op = sd_ggml_convrot_op();
|
||||
d.supports = quant_supports;
|
||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
||||
d.name = "sd.quantize_i8_convrot.h64.vulkan.v1";
|
||||
d.supports = quant64_supports;
|
||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
||||
d.name = "sd.linear_w4_convrot.vulkan.v1";
|
||||
d.custom_op = sd_ggml_w4_op();
|
||||
d.supports = w4_supports;
|
||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
||||
d.name = "sd.linear_i8.vulkan.v1";
|
||||
d.op = GGML_OP_CUSTOM;
|
||||
d.custom_op = sd_ggml_i8_op();
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#extension GL_EXT_integer_dot_product : require
|
||||
|
||||
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter {
|
||||
uint n;
|
||||
uint rows;
|
||||
uint k;
|
||||
uint has_bias;
|
||||
}
|
||||
p;
|
||||
|
||||
layout(binding = 0, std430) readonly buffer W {
|
||||
int data_w[];
|
||||
};
|
||||
layout(binding = 1, std430) readonly buffer A {
|
||||
int data_a[];
|
||||
};
|
||||
layout(binding = 2, std430) readonly buffer WS {
|
||||
float data_ws[];
|
||||
};
|
||||
layout(binding = 3, std430) readonly buffer AS {
|
||||
float data_as[];
|
||||
};
|
||||
layout(binding = 4, std430) readonly buffer B {
|
||||
float data_b[];
|
||||
};
|
||||
layout(binding = 5, std430) writeonly buffer D {
|
||||
float data_d[];
|
||||
};
|
||||
|
||||
shared int w_dec[2][16 * 64];
|
||||
shared int x_tile[2][32 * 16];
|
||||
shared int sum_p[2][16 * 8];
|
||||
|
||||
void interleave_pair(int a, int b, uint buf, uint li) {
|
||||
const uint a0 = uint(a);
|
||||
const uint b0 = uint(b);
|
||||
x_tile[buf][li + 0] = int((a0 & 0x000000FFu) | ((a0 >> 8) & 0x0000FF00u) | ((b0 << 16) & 0x00FF0000u) | ((b0 << 8) & 0xFF000000u));
|
||||
x_tile[buf][li + 1] = int(((a0 >> 8) & 0x000000FFu) | ((a0 >> 16) & 0x0000FF00u) | ((b0 << 8) & 0x00FF0000u) | (b0 & 0xFF000000u));
|
||||
}
|
||||
|
||||
void main() {
|
||||
const uint out_base = gl_WorkGroupID.x * 64u;
|
||||
const uint lx = gl_LocalInvocationID.x;
|
||||
const uint ly = gl_LocalInvocationID.y;
|
||||
const uint idx = gl_LocalInvocationIndex;
|
||||
const uint w_words = p.k / 8u;
|
||||
const uint a_words = p.k / 4u;
|
||||
const uint groups = p.k / 16u;
|
||||
|
||||
const uint dc = idx & 63u;
|
||||
const uint dw = idx >> 6u;
|
||||
const uint ocol = out_base + dc;
|
||||
|
||||
const uint fl = lx & 7u;
|
||||
const uint ft = lx >> 3u;
|
||||
const uint fr = ft == 0u ? gl_WorkGroupID.y * 32u + ly : gl_WorkGroupID.y * 32u + ly + 16u;
|
||||
const uint tr = ft * 16u + ly;
|
||||
int my_sum = 0;
|
||||
int acc0a = 0;
|
||||
int acc0b = 0;
|
||||
int acc1a = 0;
|
||||
int acc1b = 0;
|
||||
int acc2a = 0;
|
||||
int acc2b = 0;
|
||||
int acc3a = 0;
|
||||
int acc3b = 0;
|
||||
int acc4a = 0;
|
||||
int acc4b = 0;
|
||||
int acc5a = 0;
|
||||
int acc5b = 0;
|
||||
int acc6a = 0;
|
||||
int acc6b = 0;
|
||||
int acc7a = 0;
|
||||
int acc7b = 0;
|
||||
|
||||
{
|
||||
[[unroll]] for (uint j = 0; j < 2; ++j) {
|
||||
const uint pw = dw * 2 + j;
|
||||
const uint widx = pw;
|
||||
const uint wv = ocol < p.n && widx < w_words
|
||||
? uint(data_w[ocol * w_words + widx])
|
||||
: 0u;
|
||||
w_dec[0][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
|
||||
w_dec[0][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
|
||||
}
|
||||
const uint gi = fl >> 1u;
|
||||
const uint xb = gi * 4u + (fl & 1u) * 2u;
|
||||
const int a = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
|
||||
const int b = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
|
||||
interleave_pair(a, b, 0, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
|
||||
my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
|
||||
}
|
||||
barrier();
|
||||
|
||||
for (uint gs = 0; gs < groups; gs += 4) {
|
||||
const uint buf = (gs >> 2) & 1u;
|
||||
|
||||
if (gs + 4 < groups) {
|
||||
const uint nbuf = buf ^ 1u;
|
||||
const uint ngs = gs + 4;
|
||||
[[unroll]] for (uint j = 0; j < 2; ++j) {
|
||||
const uint pw = dw * 2 + j;
|
||||
const uint widx = ngs * 2 + pw;
|
||||
const uint wv = ocol < p.n && widx < w_words
|
||||
? uint(data_w[ocol * w_words + widx])
|
||||
: 0u;
|
||||
w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
|
||||
w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
|
||||
}
|
||||
const uint gi = ngs + (fl >> 1u);
|
||||
const uint xb = gi * 4u + (fl & 1u) * 2u;
|
||||
const int a = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
|
||||
const int b = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
|
||||
interleave_pair(a, b, nbuf, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
|
||||
my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
|
||||
}
|
||||
|
||||
[[unroll]] for (uint gl = 0; gl < 4; ++gl) {
|
||||
const int u0e = x_tile[buf][ly * 16u + gl * 4u + 0u];
|
||||
const int u0o = x_tile[buf][ly * 16u + gl * 4u + 1u];
|
||||
const int u0e2 = x_tile[buf][ly * 16u + gl * 4u + 2u];
|
||||
const int u0o2 = x_tile[buf][ly * 16u + gl * 4u + 3u];
|
||||
const int u1e = x_tile[buf][(16u + ly) * 16u + gl * 4u + 0u];
|
||||
const int u1o = x_tile[buf][(16u + ly) * 16u + gl * 4u + 1u];
|
||||
const int u1e2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 2u];
|
||||
const int u1o2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 3u];
|
||||
[[unroll]] for (uint jp = 0; jp < 2; ++jp) {
|
||||
const int ue0 = jp == 0u ? u0e : u0e2;
|
||||
const int uo0 = jp == 0u ? u0o : u0o2;
|
||||
const int ue1 = jp == 0u ? u1e : u1e2;
|
||||
const int uo1 = jp == 0u ? u1o : u1o2;
|
||||
const uint we = (gl * 4u + jp * 2u) * 64u;
|
||||
const uint wo = we + 64u;
|
||||
acc0a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue0));
|
||||
acc1a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue0));
|
||||
acc2a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue0));
|
||||
acc3a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue0));
|
||||
acc0b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo0));
|
||||
acc1b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo0));
|
||||
acc2b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo0));
|
||||
acc3b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo0));
|
||||
acc4a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue1));
|
||||
acc5a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue1));
|
||||
acc6a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue1));
|
||||
acc7a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue1));
|
||||
acc4b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo1));
|
||||
acc5b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo1));
|
||||
acc6b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo1));
|
||||
acc7b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo1));
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
sum_p[ft][ly * 8u + fl] = my_sum;
|
||||
barrier();
|
||||
int sum0 = 0;
|
||||
int sum1 = 0;
|
||||
[[unroll]] for (uint i = 0; i < 8; ++i) {
|
||||
sum0 += sum_p[0][ly * 8u + i];
|
||||
sum1 += sum_p[1][ly * 8u + i];
|
||||
}
|
||||
|
||||
const int accv[2][4] = {
|
||||
{acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
|
||||
{acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
|
||||
};
|
||||
const int sumv[2] = {sum0, sum1};
|
||||
[[unroll]] for (uint t = 0; t < 2u; ++t) {
|
||||
const uint rt = gl_WorkGroupID.y * 32u + ly + t * 16u;
|
||||
if (rt < p.rows) {
|
||||
[[unroll]] for (uint j = 0; j < 4; ++j) {
|
||||
const uint oj = out_base + j * 16u + lx;
|
||||
if (oj < p.n) {
|
||||
float value = float(accv[t][j] - 8 * sumv[t]) * data_as[rt] * data_ws[oj];
|
||||
if (p.has_bias != 0) {
|
||||
value += data_b[oj];
|
||||
}
|
||||
data_d[rt * p.n + oj] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
#version 460
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#extension GL_EXT_integer_dot_product : require
|
||||
|
||||
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter {
|
||||
uint n;
|
||||
uint rows;
|
||||
uint k;
|
||||
uint has_bias;
|
||||
}
|
||||
p;
|
||||
|
||||
layout(binding = 0, std430) readonly buffer W {
|
||||
int data_w[];
|
||||
};
|
||||
layout(binding = 1, std430) readonly buffer A {
|
||||
int data_a[];
|
||||
};
|
||||
layout(binding = 2, std430) readonly buffer CB {
|
||||
float data_cb[];
|
||||
};
|
||||
layout(binding = 3, std430) readonly buffer SC {
|
||||
float data_sc[];
|
||||
};
|
||||
layout(binding = 4, std430) readonly buffer SR {
|
||||
uint data_sr[];
|
||||
};
|
||||
layout(binding = 5, std430) readonly buffer AS {
|
||||
float data_as[];
|
||||
};
|
||||
layout(binding = 6, std430) readonly buffer B {
|
||||
float data_b[];
|
||||
};
|
||||
layout(binding = 7, std430) writeonly buffer D {
|
||||
float data_d[];
|
||||
};
|
||||
|
||||
shared int w_dec[2][16 * 64];
|
||||
shared int x_tile[2][32 * 16];
|
||||
shared float s_dec[2][4 * 64];
|
||||
shared int lut[256];
|
||||
shared float s_lut[256];
|
||||
shared int cb8s[16];
|
||||
shared float s_cb;
|
||||
|
||||
float f8_e4m3_to_f32(uint v) {
|
||||
const uint sign = (v >> 7) & 1u;
|
||||
const uint exp = (v >> 3) & 0xFu;
|
||||
const uint man = v & 0x7u;
|
||||
float result;
|
||||
if (exp == 0xFu && man == 0x7u) {
|
||||
result = uintBitsToFloat(0x7FC00000u);
|
||||
} else if (exp == 0u) {
|
||||
result = ldexp(float(man), -9);
|
||||
} else {
|
||||
result = ldexp(1.0 + float(man) / 8.0, int(exp) - 7);
|
||||
}
|
||||
return sign != 0u ? -result : result;
|
||||
}
|
||||
|
||||
void main() {
|
||||
const uint out_base = gl_WorkGroupID.x * 64u;
|
||||
const uint lx = gl_LocalInvocationID.x;
|
||||
const uint ly = gl_LocalInvocationID.y;
|
||||
|
||||
const uint row0 = gl_WorkGroupID.y * 32u + ly;
|
||||
const uint row1 = row0 + 16u;
|
||||
const uint idx = gl_LocalInvocationIndex;
|
||||
const uint w_words = p.k / 8u;
|
||||
const uint a_words = p.k / 4u;
|
||||
const uint groups = p.k / 16u;
|
||||
const uint sr_words = groups / 4u;
|
||||
|
||||
const uint dc = idx & 63u;
|
||||
const uint dw = idx >> 6u;
|
||||
const uint ocol = out_base + dc;
|
||||
float sum0 = 0.0;
|
||||
float sum1 = 0.0;
|
||||
float sum2 = 0.0;
|
||||
float sum3 = 0.0;
|
||||
float sum4 = 0.0;
|
||||
float sum5 = 0.0;
|
||||
float sum6 = 0.0;
|
||||
float sum7 = 0.0;
|
||||
|
||||
if (idx == 0u) {
|
||||
float amax = 0.0;
|
||||
for (uint i = 0; i < 16; ++i) {
|
||||
amax = max(amax, abs(data_cb[i]));
|
||||
}
|
||||
s_cb = amax / 127.0;
|
||||
for (uint i = 0; i < 16; ++i) {
|
||||
const int v = (s_cb == 0.0 ? 0 : int(roundEven(data_cb[i] / s_cb)));
|
||||
cb8s[i] = v < -127 ? -127 : (v > 127 ? 127 : v);
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
|
||||
lut[idx] = (cb8s[idx & 0xFu] & 0xFF) | (cb8s[idx >> 4] << 8);
|
||||
s_lut[idx] = f8_e4m3_to_f32(idx);
|
||||
barrier();
|
||||
|
||||
{
|
||||
[[unroll]] for (uint j = 0; j < 2; ++j) {
|
||||
const uint pw = dw * 2 + j;
|
||||
const uint widx = pw;
|
||||
const uint wv = ocol < p.n && widx < w_words
|
||||
? uint(data_w[ocol * w_words + widx])
|
||||
: 0u;
|
||||
|
||||
w_dec[0][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
|
||||
w_dec[0][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
|
||||
}
|
||||
{
|
||||
const uint g = dw;
|
||||
float s = 0.0;
|
||||
if (ocol < p.n && g < groups) {
|
||||
const uint sr_word = data_sr[ocol * sr_words + g / 4u];
|
||||
s = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
|
||||
}
|
||||
s_dec[0][dw * 64 + dc] = s_cb * s;
|
||||
}
|
||||
if (lx < 16u) {
|
||||
const uint aidx = lx;
|
||||
x_tile[0][ly * 16 + lx] = row0 < p.rows && aidx < a_words
|
||||
? data_a[row0 * a_words + aidx]
|
||||
: 0;
|
||||
x_tile[0][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
|
||||
? data_a[row1 * a_words + aidx]
|
||||
: 0;
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
|
||||
for (uint gs = 0; gs < groups; gs += 4) {
|
||||
const uint buf = (gs >> 2) & 1u;
|
||||
|
||||
if (gs + 4 < groups) {
|
||||
const uint nbuf = buf ^ 1u;
|
||||
const uint ngs = gs + 4;
|
||||
[[unroll]] for (uint j = 0; j < 2; ++j) {
|
||||
const uint pw = dw * 2 + j;
|
||||
const uint widx = ngs * 2 + pw;
|
||||
const uint wv = ocol < p.n && widx < w_words
|
||||
? uint(data_w[ocol * w_words + widx])
|
||||
: 0u;
|
||||
w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
|
||||
w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
|
||||
}
|
||||
{
|
||||
const uint g = ngs + dw;
|
||||
float s = 0.0;
|
||||
if (ocol < p.n && g < groups) {
|
||||
const uint sr_word = data_sr[ocol * sr_words + g / 4u];
|
||||
s = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
|
||||
}
|
||||
s_dec[nbuf][dw * 64 + dc] = s_cb * s;
|
||||
}
|
||||
if (lx < 16u) {
|
||||
const uint aidx = ngs * 4 + lx;
|
||||
x_tile[nbuf][ly * 16 + lx] = row0 < p.rows && aidx < a_words
|
||||
? data_a[row0 * a_words + aidx]
|
||||
: 0;
|
||||
x_tile[nbuf][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
|
||||
? data_a[row1 * a_words + aidx]
|
||||
: 0;
|
||||
}
|
||||
}
|
||||
|
||||
float partial0 = 0.0;
|
||||
float partial1 = 0.0;
|
||||
float partial2 = 0.0;
|
||||
float partial3 = 0.0;
|
||||
float partial4 = 0.0;
|
||||
float partial5 = 0.0;
|
||||
float partial6 = 0.0;
|
||||
float partial7 = 0.0;
|
||||
{
|
||||
[[unroll]] for (uint gl = 0; gl < 4; ++gl) {
|
||||
int accg0 = 0;
|
||||
int accg1 = 0;
|
||||
int accg2 = 0;
|
||||
int accg3 = 0;
|
||||
int accg4 = 0;
|
||||
int accg5 = 0;
|
||||
int accg6 = 0;
|
||||
int accg7 = 0;
|
||||
[[unroll]] for (uint q = 0; q < 4; ++q) {
|
||||
const int a0 = x_tile[buf][ly * 16 + gl * 4 + q];
|
||||
const int a1 = x_tile[buf][(16u + ly) * 16 + gl * 4 + q];
|
||||
accg0 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a0));
|
||||
accg1 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a0));
|
||||
accg2 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a0));
|
||||
accg3 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a0));
|
||||
accg4 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a1));
|
||||
accg5 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a1));
|
||||
accg6 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a1));
|
||||
accg7 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a1));
|
||||
}
|
||||
const float s0 = s_dec[buf][gl * 64 + lx];
|
||||
const float s1 = s_dec[buf][gl * 64 + 16u + lx];
|
||||
const float s2 = s_dec[buf][gl * 64 + 32u + lx];
|
||||
const float s3 = s_dec[buf][gl * 64 + 48u + lx];
|
||||
partial0 += float(accg0) * s0;
|
||||
partial1 += float(accg1) * s1;
|
||||
partial2 += float(accg2) * s2;
|
||||
partial3 += float(accg3) * s3;
|
||||
partial4 += float(accg4) * s0;
|
||||
partial5 += float(accg5) * s1;
|
||||
partial6 += float(accg6) * s2;
|
||||
partial7 += float(accg7) * s3;
|
||||
}
|
||||
}
|
||||
sum0 += partial0;
|
||||
sum1 += partial1;
|
||||
sum2 += partial2;
|
||||
sum3 += partial3;
|
||||
sum4 += partial4;
|
||||
sum5 += partial5;
|
||||
sum6 += partial6;
|
||||
sum7 += partial7;
|
||||
barrier();
|
||||
}
|
||||
|
||||
const float sumv[8] = {sum0, sum1, sum2, sum3, sum4, sum5, sum6, sum7};
|
||||
[[unroll]] for (uint t = 0; t < 2u; ++t) {
|
||||
const uint rt = t == 0u ? row0 : row1;
|
||||
if (rt < p.rows) {
|
||||
[[unroll]] for (uint j = 0; j < 4; ++j) {
|
||||
const uint oj = out_base + j * 16u + lx;
|
||||
if (oj < p.n) {
|
||||
float value = sumv[t * 4u + j] * data_as[rt];
|
||||
if (p.has_bias != 0) {
|
||||
value += data_b[oj];
|
||||
}
|
||||
data_d[rt * p.n + oj] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
#ifndef GROUP_SIZE
|
||||
#define GROUP_SIZE 256
|
||||
#endif
|
||||
#if GROUP_SIZE != 64 && GROUP_SIZE != 256
|
||||
#error unsupported GROUP_SIZE
|
||||
#endif
|
||||
#define GROUP_THREADS (GROUP_SIZE / 4)
|
||||
|
||||
#if GROUP_SIZE == 256
|
||||
#define HADAMARD_NORM (1.0 / 16.0)
|
||||
#else
|
||||
#define HADAMARD_NORM (1.0 / 8.0)
|
||||
#endif
|
||||
|
||||
layout(local_size_x = GROUP_THREADS, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter {
|
||||
uint k;
|
||||
uint rows;
|
||||
}
|
||||
p;
|
||||
|
||||
layout(binding = 0, std430) readonly buffer A {
|
||||
float data_a[];
|
||||
};
|
||||
layout(binding = 1, std430) writeonly buffer D {
|
||||
uint data_d[];
|
||||
};
|
||||
layout(binding = 2, std430) writeonly buffer S {
|
||||
float data_s[];
|
||||
};
|
||||
|
||||
shared float values[GROUP_SIZE];
|
||||
shared float maxima[GROUP_THREADS];
|
||||
|
||||
void transform_group(const uint row, const uint group) {
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
const uint base_offset = row * p.k + group * GROUP_SIZE;
|
||||
|
||||
[[unroll]] for (uint i = 0; i < 4; ++i) {
|
||||
const uint index = tid + i * GROUP_THREADS;
|
||||
values[index] = data_a[base_offset + index] * HADAMARD_NORM;
|
||||
}
|
||||
barrier();
|
||||
|
||||
[[unroll]] for (uint stride = 1; stride < GROUP_SIZE; stride *= 4) {
|
||||
const uint base = (tid / stride) * 4 * stride + tid % stride;
|
||||
const uint i0 = base;
|
||||
const uint i1 = i0 + stride;
|
||||
const uint i2 = i1 + stride;
|
||||
const uint i3 = i2 + stride;
|
||||
const float a = values[i0];
|
||||
const float b = values[i1];
|
||||
const float c = values[i2];
|
||||
const float d = values[i3];
|
||||
values[i0] = a + b + c - d;
|
||||
values[i1] = a + b - c + d;
|
||||
values[i2] = a - b + c + d;
|
||||
values[i3] = -a + b + c + d;
|
||||
barrier();
|
||||
}
|
||||
}
|
||||
|
||||
uint quantize_value(const float value, const float inv_scale) {
|
||||
const int quantized = clamp(int(roundEven(value * inv_scale)), -127, 127);
|
||||
return uint(quantized) & 0xffu;
|
||||
}
|
||||
|
||||
void quantize_row(const uint row) {
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
const uint groups = p.k / GROUP_SIZE;
|
||||
float local_max = 0.0;
|
||||
|
||||
for (uint group = 0; group < groups; ++group) {
|
||||
transform_group(row, group);
|
||||
[[unroll]] for (uint i = 0; i < 4; ++i) {
|
||||
local_max = max(local_max, abs(values[tid + i * GROUP_THREADS]));
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
maxima[tid] = local_max;
|
||||
barrier();
|
||||
[[unroll]] for (uint stride = GROUP_THREADS / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
maxima[tid] = max(maxima[tid], maxima[tid + stride]);
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
data_s[row] = maxima[0] / 127.0;
|
||||
}
|
||||
barrier();
|
||||
|
||||
const float row_scale = maxima[0] / 127.0;
|
||||
const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
|
||||
for (uint group = 0; group < groups; ++group) {
|
||||
transform_group(row, group);
|
||||
const uint index = tid * 4;
|
||||
const uint packed = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
|
||||
data_d[(row * p.k + group * GROUP_SIZE) / 4 + tid] = packed;
|
||||
barrier();
|
||||
}
|
||||
}
|
||||
|
||||
void main() {
|
||||
for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
|
||||
quantize_row(row);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user