Commit f5df5f8 for stable-diffusion.cpp
commit f5df5f8b84f139f1e80496902da90798a0d92896
Author: leejet <leejet714@gmail.com>
Date: Fri Nov 6 01:59:05 2026 +0800
perf: accelerate packed INT4 CUDA inference with CUTLASS (#2278)
diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md
index 49e77bc..f1e921a 100644
--- a/docs/int8_convrot.md
+++ b/docs/int8_convrot.md
@@ -123,13 +123,18 @@ Both formats imply convrot rotation even without a `convrot` flag; missing
Both formats reuse the convrot activation cache when the input scale is `1` and
produce F32 outputs with optional bias. CPU provides a portable implementation.
-CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS;
-CUDA and supported AMD targets use packed byte-dot instructions, with portable
-byte arithmetic for other targets. Vulkan requires integer dot product support
-and at least 13312 bytes of shared memory for W4A4 or 16452 bytes for W4A8.
-GPU W4 kernels require the logical input width to be divisible by 64; H64 and
-H256 are supported. Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA
-by the active compilation target's
+CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS.
+With CUTLASS enabled, NVIDIA
+SM80 and newer devices use mixed-input S8 x S4 GEMM for eligible W4A4 shapes,
+reusing the INT8 scale epilogue without expanding the full weight tensor. W4A8
+uses CUTLASS integer MMA per 16-value group with FP32 scale restoration.
+Other CUDA shapes and builds retain the packed byte-dot path. Set
+`GGML_CUDA_W4_FORCE_DP4A=1` to compare against that path. Supported AMD targets
+use packed byte-dot instructions, with portable byte arithmetic for other targets.
+Vulkan requires integer dot product support and at least 13312 bytes of shared
+memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input
+width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU
+fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's
intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout.
There is no GPU architecture allowlist: newer targets supporting either
intrinsic automatically use its matching layout. Selection happens per device
diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt
index 593445b..e2dd3f9 100644
--- a/src/core/cuda/CMakeLists.txt
+++ b/src/core/cuda/CMakeLists.txt
@@ -27,6 +27,12 @@ if(SD_CUDA)
set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
if(SD_CUDA_CUTLASS_PATH)
+ set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS)
+ set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include")
+ set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr)
+ if(MSVC)
+ set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+ endif()
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS)
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu
index 69f91d2..9d4d92e 100644
--- a/src/core/cuda/i8_gemm.cu
+++ b/src/core/cuda/i8_gemm.cu
@@ -143,8 +143,9 @@ namespace sd_i8_cutlass {
dst[i] = src[i / 2 + (i % 2) * (n / 2)];
}
- template <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;
diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h
index 189df6e..5e4ac9f 100644
--- a/src/core/cuda/i8_gemm.h
+++ b/src/core/cuda/i8_gemm.h
@@ -3,7 +3,7 @@
#include <cuda_runtime.h>
#include "ggml-ext.h"
struct SDI8CudaGemmPlan;
-SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace);
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false);
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan);
#endif
diff --git a/src/core/cuda/w4.cu b/src/core/cuda/w4.cu
index ef46d02..98f810e 100644
--- a/src/core/cuda/w4.cu
+++ b/src/core/cuda/w4.cu
@@ -1,6 +1,13 @@
#include <climits>
#include "core/ggml_extensions.h"
#include "runtime.h"
+#ifdef SD_CUDA_W4_CUTLASS
+#include <cstdlib>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+#include "w4_mma.h"
+#endif
namespace SD_CUDA_NAMESPACE {
namespace sd_w4_cuda {
@@ -789,7 +796,37 @@ namespace SD_CUDA_NAMESPACE {
return false;
return true;
}
- static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+ struct Plan {
+ bool mma = false;
+ size_t gemm_offset = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+ SDI8CudaGemmPlan* gemm = nullptr;
+ ~Plan() { sd_cuda_i8_gemm_free(gemm); }
+#endif
+ };
+ static bool use_mma(ggml_backend_t backend, const SDW4Operation& op) {
+#ifdef SD_CUDA_W4_CUTLASS
+ if (!aligned(op.input, 16) || !aligned(op.weight, 16))
+ return false;
+ static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_W4_FORCE_DP4A"); return v && std::atoi(v); }();
+ if (disabled)
+ return false;
+ static std::mutex mutex;
+ static std::unordered_map<int, bool> devices;
+ const auto device = cuda_device_info(ggml_backend_get_device(backend)).device;
+ std::lock_guard<std::mutex> lock(mutex);
+ auto found = devices.find(device);
+ if (found != devices.end())
+ return found->second;
+ cudaDeviceProp prop = {};
+ const bool available = cudaGetDeviceProperties(&prop, device) == cudaSuccess && prop.major >= 8;
+ devices.emplace(device, available);
+ return available;
+#else
+ return false;
+#endif
+ }
+ static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int, void** state, size_t* bytes, void*) {
SDW4Operation op;
if (!sd_ggml_w4_describe(nodes[0], &op))
return GGML_STATUS_FAILED;
@@ -812,7 +849,22 @@ namespace SD_CUDA_NAMESPACE {
may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst));
}
}
- *bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 ? 1 : 0) + (may_alias ? 1 : 0));
+ auto* plan = new Plan;
+ plan->mma = use_mma(backend, op);
+ *state = plan;
+ size_t gemm_bytes = 0;
+#ifdef SD_CUDA_W4_CUTLASS
+ if (plan->mma && op.kind == SDW4Kind::W4A4) {
+ if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16))
+ plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true);
+ plan->mma = plan->gemm != nullptr;
+ if (plan->gemm && gemm_bytes == 0)
+ gemm_bytes = 256;
+ }
+#endif
+ const size_t tensor_bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 && !plan->mma ? 1 : 0) + (may_alias ? 1 : 0));
+ plan->gemm_offset = (tensor_bytes + 255) / 256 * 256;
+ *bytes = plan->gemm_offset + gemm_bytes;
return GGML_STATUS_SUCCESS;
}
static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
@@ -821,7 +873,7 @@ namespace SD_CUDA_NAMESPACE {
const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
}
- static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) {
+ static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void* state, void*) {
auto* dst = nodes[0];
SDW4Operation op;
ggml_ext_gpu_context context;
@@ -840,7 +892,21 @@ namespace SD_CUDA_NAMESPACE {
const auto* ws = static_cast<const float*>(op.weight_scale->data);
const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16);
- if (op.kind == SDW4Kind::W4A4) {
+#ifdef SD_CUDA_W4_CUTLASS
+ auto* plan = static_cast<Plan*>(state);
+ if (plan->mma) {
+ if (plan->gemm) {
+ auto* workspace = scratch ? reinterpret_cast<char*>(scratch) + plan->gemm_offset : nullptr;
+ if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast<const int8_t*>(x), reinterpret_cast<const int8_t*>(w), xs, ws, output, workspace, stream) != GGML_STATUS_SUCCESS)
+ return GGML_STATUS_FAILED;
+ } else {
+ const dim3 tiles(unsigned((n + 63) / 64), unsigned((rows + 127) / 128));
+ sd_w4_mma::w4a8<128, 64, 128><<<tiles, 256, 0, stream>>>(reinterpret_cast<const uint8_t*>(w), reinterpret_cast<const int8_t*>(x), xs, ws,
+ static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data), bias, output, int(rows), int(n), int(k));
+ }
+ } else
+#endif
+ if (op.kind == SDW4Kind::W4A4) {
auto* accum = reinterpret_cast<int32_t*>(scratch + (alias ? ggml_nelements(dst) : 0));
mul_mat_w4a4_convrot_dp4a<<<grid, block, 0, stream>>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2));
if (cuda_last_error() != cuda_success)
@@ -867,7 +933,7 @@ namespace SD_CUDA_NAMESPACE {
d.supports = sd_w4_cuda::supports;
d.prepare = sd_w4_cuda::prepare;
d.execute = sd_w4_cuda::execute;
- d.destroy = [](void*, void*) {};
+ d.destroy = [](void* state, void*) { delete static_cast<sd_w4_cuda::Plan*>(state); };
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
diff --git a/src/core/cuda/w4_mma.h b/src/core/cuda/w4_mma.h
new file mode 100644
index 0000000..94629bf
--- /dev/null
+++ b/src/core/cuda/w4_mma.h
@@ -0,0 +1,138 @@
+#ifndef __SD_CORE_CUDA_W4_MMA_H__
+#define __SD_CORE_CUDA_W4_MMA_H__
+
+#include <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