feat: support packed INT4 convrot models (#2276)

Co-authored-by: noctrex <noctrex@gmail.com>
This commit is contained in:
leejetandnoctrex authored and GitHub committed 2026-11-06 00:35:58 +08:00
1 parent 58bb3c1b1c
commit 7952693004
48 files changed
+1992 -91

No files matched your search

+10
View File
@@ -104,6 +104,16 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions.
## Backend services
Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by
`ggml_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the
explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization
node during graph construction. Logical output dimensions are stored in the
operator attributes, so backend copies of the packed activation do not need to
retain the producer's op marker. W4A4 uses row scales; W4A8 additionally exposes
the codebook, channel scales, and FP8 relative scales as tensor dependencies.
CPU, CUDA/HIP/MUSA, and Vulkan register kernels through the same extension API
as INT8. No ggml submodule changes are required.
Services are discovered through `ggml_backend_reg_get_proc_address`:
| Service | Purpose |
+40 -1
View File
@@ -100,7 +100,46 @@ The packed runtime activation tensor contains the I8 activation rows and their f
LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
The shared CUDA/HIP/MUSA convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution.
The shared CUDA/HIP/MUSA convrot activation paths support H64 and H256. Vulkan also supports H64 activation quantization; its INT8 matrix multiplication path currently requires H256.
## Packed INT4 convrot
ComfyUI `convrot_w4a4` and `asym_w4a8_int8` safetensors load directly, including
checkpoints that mix packed INT4, INT8 tensorwise, and floating-point layers.
Packed weights retain their checkpoint byte layout; they are not expanded or
converted to GGUF Q4 blocks at load time.
- `convrot_w4a4`: signed 4-bit weights, low nibble first, with F32 per-output
`weight_scale`. Activations use the same INT8 convrot quantization as above;
the format name does not select a native INT4 activation kernel.
- `asym_w4a8_int8`: 4-bit indices into `weight_codebook` (16 F32 values), F32
`weight_s_channel`, and F8_E4M3 `weight_s_rel` per 16 input features. The
metadata must specify `group_size: 16`. CPU evaluates the floating-point
codebook; GPU kernels requantize it to INT8 with `amax / 127`, introducing a
small additional rounding error.
Both formats imply convrot rotation even without a `convrot` flag; missing
`convrot_groupsize` defaults to 256.
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
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
image, including multi-architecture builds; targets supporting neither intrinsic
retain the integer-dot fallback. W4A4 retains integer dot products on HIP.
LoRA application remains at runtime, using the original unrotated activations
and logical weight dimensions. Packed weights and their codebook/scale tensors
are excluded from weight type conversion. `SD_USE_UPSTREAM_GGML` builds reject
these formats because they require the application extension API.
## Example
+2
View File
@@ -60,6 +60,8 @@ static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage
```
Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes.
Use `logical_ne0()` for Linear input widths; packed INT4 weights store half of
the logical width in `ne[0]`. Keep raw `ne` values for storage layout calculations.
Do not load or parse tensor data for config detection.
Detection should respect `prefix`. For nested weights, construct full names from
+1
View File
@@ -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);
}
+1
View File
@@ -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
+95
View File
@@ -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
+3 -1
View File
@@ -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()
+64 -4
View File
@@ -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,6 +501,9 @@ 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;
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);
@@ -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;
+2
View File
@@ -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);
+873
View File
@@ -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));
}
}
+95
View File
@@ -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;
+19
View File
@@ -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,
+3
View File
@@ -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)
+70 -2
View File
@@ -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,8 +554,11 @@ 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;
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);
@@ -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);
}
}
+2 -2
View File
@@ -141,7 +141,7 @@ namespace IPAdapter {
}
auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
if (proj_in_iter != tensor_storage_map.end()) {
embed_dim = proj_in_iter->second.ne[0];
embed_dim = proj_in_iter->second.logical_ne0();
}
auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
if (proj_out_iter != tensor_storage_map.end()) {
@@ -168,7 +168,7 @@ namespace IPAdapter {
}
auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
if (proj_iter != tensor_storage_map.end()) {
clip_dim = proj_iter->second.ne[0];
clip_dim = proj_iter->second.logical_ne0();
out_dim = proj_iter->second.ne[1];
}
num_tokens = out_dim / ctx_dim;
+10
View File
@@ -1127,6 +1127,16 @@ public:
ggml_tensor* output,
const std::string& prefix,
WeightAdapter::ForwardParams forward_params) override {
// Packed INT4 stores half the logical input width. Adapter shape checks
// and reshapes need the logical weight dimensions, without decoding it.
ggml_tensor logical_weight = *w;
if (forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_LINEAR &&
w->type == GGML_TYPE_I8 && x->ne[0] == 2 * w->ne[0]) {
logical_weight.ne[0] = x->ne[0];
for (int i = 1; i < GGML_MAX_DIMS; ++i)
logical_weight.nb[i] = logical_weight.nb[i - 1] * logical_weight.ne[i - 1];
w = &logical_weight;
}
for (auto& lora_model : lora_models) {
ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
if (weight_diff != nullptr) {
+1 -1
View File
@@ -319,7 +319,7 @@ protected:
tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
has_ip = true;
int64_t inner_dim = d_head * n_head;
int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").logical_ne0();
blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
}
+36 -5
View File
@@ -148,6 +148,7 @@ protected:
bool force_f32;
bool force_prec_f32;
bool has_weight_scale = false;
int w4_convrot_kind = 0;
bool int8_convrot = false;
int int8_convrot_group_size = 0;
float scale;
@@ -158,16 +159,21 @@ protected:
has_weight_scale = false;
int8_convrot = false;
int8_convrot_group_size = 0;
auto weight_storage = tensor_storage_map.find(prefix + "weight");
w4_convrot_kind = weight_storage != tensor_storage_map.end() ? weight_storage->second.w4_convrot_kind : 0;
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
if (!w4_convrot_kind && (in_features % ggml_blck_size(wtype) != 0 || force_f32)) {
wtype = GGML_TYPE_F32;
}
params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
if (w4_convrot_kind) {
GGML_ASSERT(wtype == GGML_TYPE_I8 && in_features % 2 == 0);
GGML_ASSERT(weight_storage->second.ne[0] == in_features / 2 && weight_storage->second.ne[1] == out_features);
}
params["weight"] = ggml_new_tensor_2d(ctx, wtype, w4_convrot_kind ? in_features / 2 : in_features, out_features);
if (bias) {
enum ggml_type wtype = GGML_TYPE_F32;
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features);
}
auto weight_storage = tensor_storage_map.find(prefix + "weight");
const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale");
if (weight_scale_storage != tensor_storage_map.end()) {
@@ -176,9 +182,16 @@ protected:
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
has_weight_scale = true;
}
if (is_int8_tensorwise) {
if (w4_convrot_kind == 2) {
params["weight_codebook"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16);
params["weight_s_channel"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
#ifndef SD_USE_UPSTREAM_GGML
params["weight_s_rel"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F8_E4M3, in_features / 16, out_features);
#endif
}
if (is_int8_tensorwise || w4_convrot_kind) {
GGML_ASSERT(wtype == GGML_TYPE_I8);
GGML_ASSERT(has_weight_scale);
GGML_ASSERT(w4_convrot_kind == 2 || has_weight_scale);
int8_convrot = weight_storage->second.int8_convrot;
int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
}
@@ -258,6 +271,24 @@ public:
x = cached->second;
}
}
#endif
#ifndef SD_USE_UPSTREAM_GGML
if (w4_convrot_kind) {
if (scale != 1.f)
x = ggml_ext_scale(ctx->ggml_ctx, x, scale);
out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x,
w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
scale == 1.f ? b : nullptr, static_cast<SDW4Kind>(w4_convrot_kind),
int8_convrot ? int8_convrot_group_size : 0);
GGML_ASSERT(out);
if (scale != 1.f) {
out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale);
if (b)
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
} else
#endif
out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
x,
+2 -2
View File
@@ -72,10 +72,10 @@ namespace Boogu {
}
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.in_channels = tensor_storage.ne[0] / patch_area;
config.in_channels = tensor_storage.logical_ne0() / patch_area;
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.instruction_feat_dim = tensor_storage.ne[0];
config.instruction_feat_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
detected_head_dim = tensor_storage.ne[0];
+1 -1
View File
@@ -42,7 +42,7 @@ namespace ErnieImage {
config.in_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
config.text_in_dim = tensor_storage.ne[0];
config.text_in_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
detected_head_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
+2 -2
View File
@@ -31,12 +31,12 @@ namespace Fibo2 {
}
auto key = name.substr(prefix.size() + 1);
if (key == "x_embedder.weight") {
config.channels = tensor.ne[0];
config.channels = tensor.logical_ne0();
config.dim = tensor.ne[1];
} else if (key == "perceiver.latent_queries") {
config.max_gist = tensor.ne[1];
} else if (key == "perceiver.layers.0.cross_attention.to_k.weight" || key == "perceiver.layers.0.cross_attention.to_kv.weight") {
config.text_dim = tensor.ne[0];
config.text_dim = tensor.logical_ne0();
config.fused_kv = key == "perceiver.layers.0.cross_attention.to_kv.weight";
} else if (key == "layers.0.attention.norm_q.weight") {
config.head_dim = tensor.ne[0];
+1 -1
View File
@@ -158,7 +158,7 @@ namespace Flux {
}
}
if (ends_with(name, "txt_in.weight")) {
config.context_in_dim = tensor_storage.ne[0];
config.context_in_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
}
if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) {
+3 -3
View File
@@ -27,12 +27,12 @@ namespace Flux3 {
std::string key = name.substr(prefix.size() + 1);
if (key == "txt_in.weight") {
detected = true;
config.context_dim = tensor.ne[0];
config.context_dim = tensor.logical_ne0();
config.hidden_size = tensor.ne[1];
} else if (key == "emb_in.image.weight") {
config.channels = tensor.ne[0];
config.channels = tensor.logical_ne0();
} else if (key == "single_blocks.0.mlp_out.weight") {
config.mlp_dim = tensor.ne[0];
config.mlp_dim = tensor.logical_ne0();
} else if (key == "ref_time_in.in_layer.weight") {
config.flow_map = true;
}
+5 -5
View File
@@ -92,11 +92,11 @@ namespace Krea2 {
continue;
}
if (ends_with(name, "first.weight") && tensor_storage.n_dims == 2) {
config.in_channels = tensor_storage.ne[0] / (config.patch_size * config.patch_size);
config.in_channels = tensor_storage.logical_ne0() / (config.patch_size * config.patch_size);
config.out_channels = config.in_channels;
config.features = tensor_storage.ne[1];
} else if (ends_with(name, "blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
detected_head_dim = tensor_storage.ne[0];
detected_head_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
if (detected_head_dim > 0) {
config.heads = tensor_storage.ne[1] / detected_head_dim;
@@ -106,11 +106,11 @@ namespace Krea2 {
config.kv_heads = tensor_storage.ne[1] / detected_head_dim;
}
} else if (ends_with(name, "txtfusion.projector.weight") && tensor_storage.n_dims == 2) {
config.text_layers = tensor_storage.ne[0];
config.text_layers = tensor_storage.logical_ne0();
} else if (ends_with(name, "txtfusion.layerwise_blocks.0.prenorm.scale") && tensor_storage.n_dims == 1) {
config.text_dim = tensor_storage.ne[0];
config.text_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
detected_text_head_dim = tensor_storage.ne[0];
detected_text_head_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
if (detected_text_head_dim > 0) {
config.text_heads = tensor_storage.ne[1] / detected_text_head_dim;
+2 -2
View File
@@ -36,13 +36,13 @@ namespace Lens {
continue;
}
if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
config.in_channels = tensor_storage.ne[0];
config.in_channels = tensor_storage.logical_ne0();
int64_t inner_dim = tensor_storage.ne[1];
if (config.attention_head_dim > 0) {
config.num_attention_heads = inner_dim / config.attention_head_dim;
}
} else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
config.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / config.joint_attention_dim);
config.selected_layer_count = static_cast<int>(tensor_storage.logical_ne0() / config.joint_attention_dim);
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = tensor_storage.ne[1] / patch_area;
+3 -3
View File
@@ -63,16 +63,16 @@ namespace LingBotVideo {
}
if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_dim = tensor_storage.ne[0];
int64_t patch_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
if (patch_dim % patch_volume == 0) {
config.in_channels = patch_dim / patch_volume;
}
} else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.text_dim = tensor_storage.ne[0];
config.text_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.freq_dim = tensor_storage.ne[0];
config.freq_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
int64_t head_dim = tensor_storage.ne[0];
if (head_dim > 0) {
+3 -3
View File
@@ -74,13 +74,13 @@ namespace LLaDAImage {
}
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.in_channels = tensor_storage.ne[0] / patch_area;
config.in_channels = tensor_storage.logical_ne0() / patch_area;
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.cap_feat_dim = tensor_storage.ne[0];
config.cap_feat_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.semantic_feat_dim = tensor_storage.ne[0];
config.semantic_feat_dim = tensor_storage.logical_ne0();
} else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
detected_q_dim = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {
+4 -4
View File
@@ -168,7 +168,7 @@ namespace LTXV {
LTXAVConfig config;
auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
if (patchify_proj_iter != tensor_storage_map.end()) {
config.in_channels = patchify_proj_iter->second.ne[0];
config.in_channels = patchify_proj_iter->second.logical_ne0();
config.hidden_size = patchify_proj_iter->second.ne[1];
int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
auto attn_layout = infer_attention_layout(config.hidden_size, video_heads);
@@ -178,7 +178,7 @@ namespace LTXV {
auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
if (audio_patchify_proj_iter != tensor_storage_map.end()) {
config.audio_in_channels = audio_patchify_proj_iter->second.ne[0];
config.audio_in_channels = audio_patchify_proj_iter->second.logical_ne0();
config.audio_hidden_size = audio_patchify_proj_iter->second.ne[1];
config.audio_out_channels = config.audio_in_channels;
int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
@@ -198,11 +198,11 @@ namespace LTXV {
auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
if (attn2_iter != tensor_storage_map.end()) {
config.cross_attention_dim = attn2_iter->second.ne[0];
config.cross_attention_dim = attn2_iter->second.logical_ne0();
}
auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
if (audio_attn2_iter != tensor_storage_map.end()) {
config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
config.audio_cross_attention_dim = audio_attn2_iter->second.logical_ne0();
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
config.cross_attention_adaln = true;
+5 -5
View File
@@ -73,16 +73,16 @@ namespace MiniMaxH3 {
};
if (const auto* weight = find("video_patch_proj.weight")) {
config.video_latent_channels = weight->ne[0] / 4;
config.video_latent_channels = weight->logical_ne0() / 4;
config.hidden_size = weight->ne[1];
}
if (const auto* weight = find("audio_patch_proj.weight")) {
config.audio_latent_channels = weight->ne[0];
config.audio_latent_channels = weight->logical_ne0();
}
config.num_layers = count_blocks(tensors, prefix + ".blocks.");
config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks.");
if (const auto* weight = find("blocks.0.attn.q_norm.weight")) {
config.attention_head_dim = weight->ne[0];
config.attention_head_dim = weight->logical_ne0();
}
if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) {
config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim);
@@ -91,14 +91,14 @@ namespace MiniMaxH3 {
config.ffn_hidden_size = weight->ne[1] / 2;
}
if (const auto* weight = find("condition_proj.weight")) {
config.text_dim = weight->ne[0];
config.text_dim = weight->logical_ne0();
}
if (const auto* table = find("adaln_t_table")) {
config.time_embed_dim = table->ne[0];
config.adaln_curve_grid = table->ne[1];
} else {
if (const auto* weight = find("time_embedder.proj_in.weight")) {
config.timestep_input_dim = weight->ne[0];
config.timestep_input_dim = weight->logical_ne0();
config.time_embed_hidden_size = weight->ne[1];
}
if (const auto* weight = find("time_embedder.proj_out.weight")) {
+3 -3
View File
@@ -59,14 +59,14 @@ namespace MiniT2I {
config.pca_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
config.txt_input_size = tensor_storage.ne[0];
config.txt_input_size = tensor_storage.logical_ne0();
config.txt_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
config.cond_vec_size = tensor_storage.ne[1];
} else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
int64_t inner3 = tensor_storage.ne[1];
int64_t inner = inner3 / 3;
config.hidden_size = tensor_storage.ne[0];
config.hidden_size = tensor_storage.logical_ne0();
if (config.hidden_size == 768) {
config.num_heads = 12;
config.head_dim = 64;
@@ -79,7 +79,7 @@ namespace MiniT2I {
}
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.hidden_size = tensor_storage.ne[0];
config.hidden_size = tensor_storage.logical_ne0();
config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
} else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
config.prompt_length = tensor_storage.ne[1];
+3 -3
View File
@@ -61,17 +61,17 @@ struct MMDiTConfig {
} else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_hidden_size = true;
config.adm_in_channels = tensor_storage.ne[0];
config.adm_in_channels = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
} else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_context_embed = true;
config.context_size = tensor_storage.ne[0];
config.context_size = tensor_storage.logical_ne0();
config.context_embedder_out_dim = tensor_storage.ne[1];
} else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_hidden_size = true;
config.hidden_size = tensor_storage.ne[0];
config.hidden_size = tensor_storage.logical_ne0();
int64_t patch_area = static_cast<int64_t>(config.patch_size) * config.patch_size;
if (patch_area > 0) {
config.out_channels = tensor_storage.ne[1] / patch_area;
+4 -4
View File
@@ -32,16 +32,16 @@ namespace MOSSImage {
}
auto key = name.substr(prefix.size() + 1);
if (key == "video_embedder.weight") {
config.channels = tensor.ne[0];
config.channels = tensor.logical_ne0();
config.dim = tensor.ne[1];
} else if (key == "text_embedder.proj.weight") {
config.text_dim = tensor.ne[0];
config.text_dim = tensor.logical_ne0();
} else if (key == "layers.0.proj.ffn_down.weight") {
config.ffn_dim = tensor.ne[0];
config.ffn_dim = tensor.logical_ne0();
} else if (key == "layers.0.q_norm.weight") {
config.head_dim = tensor.ne[0];
} else if (key == "time_embedding.0.weight") {
config.freq_dim = tensor.ne[0];
config.freq_dim = tensor.logical_ne0();
}
if (starts_with(key, "layers.")) {
depth = std::max<int64_t>(depth, std::stoi(key.substr(7)) + 1);
+2 -2
View File
@@ -51,10 +51,10 @@ namespace PixArt {
config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
}
if (auto w = find("y_embedder.y_proj.fc1.weight")) {
config.caption_channels = w->ne[0];
config.caption_channels = w->logical_ne0();
}
if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) {
config.cross_attention_dim = w->ne[0];
config.cross_attention_dim = w->logical_ne0();
}
if (auto w = find("blocks.0.mlp.fc1.weight")) {
config.ffn_dim = w->ne[1];
+2 -2
View File
@@ -24,14 +24,14 @@ namespace Qwen {
return it == weights.end() ? nullptr : &it->second;
};
if (auto w = find("img_in.weight")) {
config.in_channels = w->ne[0];
config.in_channels = w->logical_ne0();
config.hidden_size = w->ne[1];
}
if (auto w = find("proj_out.weight")) {
config.out_channels = w->ne[1];
}
if (auto w = find("txt_in.in_layer.weight")) {
config.context_dim = w->ne[0];
config.context_dim = w->logical_ne0();
}
if (auto w = find("transformer_blocks.0.attn.norm_q.weight")) {
config.head_dim = w->ne[0];
+1 -1
View File
@@ -31,7 +31,7 @@ namespace SefiImage {
continue;
}
if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.timestep_guidance_in_dim = tensor_storage.ne[0];
config.timestep_guidance_in_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1] * 2;
}
}
+3 -3
View File
@@ -104,13 +104,13 @@ struct UNetConfig {
}
if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) {
if (time_embed->n_dims == 2) {
config.model_channels = static_cast<int>(time_embed->ne[0]);
config.model_channels = static_cast<int>(time_embed->logical_ne0());
config.time_embed_dim = static_cast<int>(time_embed->ne[1]);
}
}
if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) {
if (label_emb->n_dims == 2) {
config.adm_in_channels = static_cast<int>(label_emb->ne[0]);
config.adm_in_channels = static_cast<int>(label_emb->logical_ne0());
config.time_embed_dim = static_cast<int>(label_emb->ne[1]);
}
}
@@ -124,7 +124,7 @@ struct UNetConfig {
continue;
}
if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) {
config.context_dim = static_cast<int>(tensor_storage.ne[0]);
config.context_dim = static_cast<int>(tensor_storage.logical_ne0());
break;
}
}
+2 -2
View File
@@ -62,10 +62,10 @@ namespace ZImage {
}
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.in_channels = tensor_storage.ne[0] / patch_area;
config.in_channels = tensor_storage.logical_ne0() / patch_area;
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.cap_feat_dim = tensor_storage.ne[0];
config.cap_feat_dim = tensor_storage.logical_ne0();
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) {
detected_head_dim = tensor_storage.ne[0];
+1 -1
View File
@@ -18,7 +18,7 @@ namespace ZImageL2P {
config.patch_size = 16;
auto x_embedder = tensors.find(prefix + ".x_embedder.weight");
if (x_embedder != tensors.end()) {
config.patch_size = static_cast<int>(std::lround(std::sqrt(static_cast<double>(x_embedder->second.ne[0] / config.in_channels))));
config.patch_size = static_cast<int>(std::lround(std::sqrt(static_cast<double>(x_embedder->second.logical_ne0() / config.in_channels))));
}
// L2P ships split to_q/to_k/to_v, which name conversion maps to qkv.weight plus
// .1/.2 parts; the base detection only sees the q part and undercounts kv heads.
+2 -2
View File
@@ -57,7 +57,7 @@ struct T5Config {
}
if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) {
if (q->n_dims == 2) {
config.model_dim = q->ne[0];
config.model_dim = q->logical_ne0();
int64_t inner_dim = q->ne[1];
config.inner_dim = inner_dim;
// Flan-T5/T5 uses d_kv=64 for common sizes.
@@ -68,7 +68,7 @@ struct T5Config {
}
if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) {
if (wi->n_dims == 2) {
config.model_dim = wi->ne[0];
config.model_dim = wi->logical_ne0();
config.ff_dim = wi->ne[1];
}
}
+1 -1
View File
@@ -75,7 +75,7 @@ namespace H3VUpsampler {
const TensorStorage* embed_mid = find_weight("embed.2.weight");
if (embed_mid != nullptr && embed_mid->n_dims >= 1) {
config.embed_dim = static_cast<int>(embed_mid->ne[0]);
config.embed_dim = static_cast<int>(embed_mid->logical_ne0());
}
std::set<int> in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx;
+45 -9
View File
@@ -111,6 +111,7 @@ struct ComfyQuantConfig {
std::string format;
bool convrot = false;
int group_size = 0;
int weight_group_size = 0;
};
static bool read_comfy_quant_config(std::ifstream& file,
@@ -140,6 +141,12 @@ static bool read_comfy_quant_config(std::ifstream& file,
config.format = json.value("format", "");
config.convrot = json.value("convrot", false);
config.group_size = json.value("convrot_groupsize", 0);
config.weight_group_size = json.value("group_size", 0);
if (config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") {
// Packed W4 formats imply rotation, including exports without the INT8 convrot flag.
config.convrot = true;
config.group_size = json.value("convrot_groupsize", 256);
}
} catch (const std::exception&) {
set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
return false;
@@ -238,8 +245,8 @@ bool read_safetensors_file(const std::string& file_path,
return false;
}
#ifdef SD_USE_UPSTREAM_GGML
if (config.format == "int8_tensorwise") {
set_error(error, "INT8 tensorwise/convrot is not supported by this ggml build (tensor '" + name + "')");
if (config.format == "int8_tensorwise" || config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") {
set_error(error, "ComfyUI INT8/packed INT4 is not supported by this ggml build (tensor '" + name + "')");
return false;
}
#endif
@@ -340,29 +347,58 @@ bool read_safetensors_file(const std::string& file_path,
if (ends_with(name, ".weight")) {
const std::string module_name = name.substr(0, name.size() - std::string(".weight").size());
auto config = comfy_quant_configs.find(module_name);
if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise") {
if (config != comfy_quant_configs.end() &&
(config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8")) {
const int w4_kind = config->second.format == "convrot_w4a4" ? 1 : config->second.format == "asym_w4a8_int8" ? 2
: 0;
if (type != GGML_TYPE_I8) {
set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'");
set_error(error, "ComfyUI quantized weight is not I8: '" + name + "'");
return false;
}
if (w4_kind && (tensor_storage.n_dims != 2 || tensor_storage.ne[0] <= 0 ||
tensor_storage.ne[0] > INT64_MAX / 2 || tensor_storage.ne[1] <= 0)) {
set_error(error, "invalid packed INT4 weight shape: '" + name + "'");
return false;
}
const int64_t input_features = tensor_storage.ne[0] * (w4_kind ? 2 : 1);
if (config->second.convrot) {
int group_size_remainder = config->second.group_size;
while (group_size_remainder > 1 && group_size_remainder % 4 == 0) {
group_size_remainder /= 4;
}
if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) {
if (group_size_remainder != 1 || config->second.group_size > 256 || input_features % config->second.group_size != 0) {
set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'");
return false;
}
}
tensor_storage.is_int8_tensorwise = true;
if (w4_kind) {
auto matches = [&](const std::string& suffix, const std::string& expected_dtype, const std::vector<int64_t>& expected_shape) {
auto entry = header_.find(module_name + suffix);
return entry != header_.end() && entry->value("dtype", "") == expected_dtype &&
(*entry)["shape"] == expected_shape;
};
const int64_t n = tensor_storage.ne[1];
const bool valid = w4_kind == 1
? (matches(".weight_scale", "F32", {n}) || matches(".weight_scale", "F32", {n, 1}))
: config->second.weight_group_size == 16 && input_features % 16 == 0 &&
matches(".weight_codebook", "F32", {16}) &&
(matches(".weight_s_channel", "F32", {n}) || matches(".weight_s_channel", "F32", {n, 1})) &&
matches(".weight_s_rel", "F8_E4M3", {n, input_features / 16});
if (!valid) {
set_error(error, "invalid packed INT4 scales/codebook for tensor '" + name + "'");
return false;
}
}
tensor_storage.is_int8_tensorwise = w4_kind == 0;
tensor_storage.w4_convrot_kind = w4_kind;
tensor_storage.int8_convrot = config->second.convrot;
tensor_storage.int8_convrot_group_size = config->second.group_size;
}
} else if (ends_with(name, ".weight_scale")) {
const std::string module_name = name.substr(0, name.size() - std::string(".weight_scale").size());
} else if (ends_with(name, ".weight_scale") || ends_with(name, ".weight_s_channel")) {
const std::string module_name = name.substr(0, name.rfind('.'));
auto config = comfy_quant_configs.find(module_name);
if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise" &&
if (config != comfy_quant_configs.end() &&
(config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8") &&
tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) {
tensor_storage.ne[0] = tensor_storage.ne[1];
tensor_storage.ne[1] = 1;
+5
View File
@@ -22,6 +22,7 @@ struct TensorStorage {
bool is_f64 = false;
bool is_i64 = false;
bool is_int8_tensorwise = false;
int w4_convrot_kind = 0;
bool int8_convrot = false;
int int8_convrot_group_size = 0;
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
@@ -51,6 +52,10 @@ struct TensorStorage {
return n;
}
int64_t logical_ne0() const {
return ne[0] * (w4_convrot_kind ? 2 : 1);
}
int64_t nbytes() const {
return nelements() * ggml_type_size(type) / ggml_blck_size(type);
}
+9 -8
View File
@@ -518,7 +518,7 @@ SDVersion ModelLoader::get_sd_version() const {
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
if (img_in != tensor_storage_map.end() && img_in->second.logical_ne0() == 128) {
return VERSION_MAGE_FLOW;
}
if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
@@ -593,7 +593,7 @@ SDVersion ModelLoader::get_sd_version() const {
}
}
auto connector = tensor_storage_map.find("embeddings_connectors.proj.weight");
if (connector != tensor_storage_map.end() && connector->second.ne[0] == 12800 &&
if (connector != tensor_storage_map.end() && connector->second.logical_ne0() == 12800 &&
connector->second.ne[1] == 5120 && tensor_storage_map.count("embeddings_connectors.sink_embedding")) {
return VERSION_P_H3_IMAGE;
}
@@ -660,7 +660,7 @@ SDVersion ModelLoader::get_sd_version() const {
tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) {
has_output_block_71 = true;
if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) {
if (tensor_storage.ne[0] == 1024)
if (tensor_storage.logical_ne0() == 1024)
has_attn_1024 = true;
}
}
@@ -721,16 +721,16 @@ SDVersion ModelLoader::get_sd_version() const {
if (has_distilled_guidance) {
return VERSION_CHROMA;
}
if (context_ebedding_weight.ne[0] == 3584) {
if (context_ebedding_weight.logical_ne0() == 3584) {
return VERSION_LONGCAT;
} else {
if (input_block_weight.ne[0] == 384) {
if (input_block_weight.logical_ne0() == 384) {
return VERSION_FLUX_FILL;
}
if (input_block_weight.ne[0] == 128) {
if (input_block_weight.logical_ne0() == 128) {
return VERSION_FLUX_CONTROLS;
}
if (input_block_weight.ne[0] == 196) {
if (input_block_weight.logical_ne0() == 196) {
return VERSION_FLEX_2;
}
return VERSION_FLUX;
@@ -1651,7 +1651,8 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const {
const std::string& name = tensor_storage.name;
if (tensor_storage.is_int8_tensorwise) {
if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind != 0 ||
ends_with(name, ".weight_codebook") || ends_with(name, ".weight_s_channel") || ends_with(name, ".weight_s_rel")) {
return false;
}
if (type != GGML_TYPE_COUNT) {
+1 -1
View File
@@ -11,7 +11,7 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot &&
a.is_int8_tensorwise == b.is_int8_tensorwise && a.w4_convrot_kind == b.w4_convrot_kind && a.int8_convrot == b.int8_convrot &&
a.int8_convrot_group_size == b.int8_convrot_group_size;
}
+2 -2
View File
@@ -1019,7 +1019,7 @@ void StableDiffusionGGML::configure_weight_loading() {
bool have_int8_tensorwise = false;
bool have_quant_type_rule = false;
for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) {
if (tensor_storage.is_int8_tensorwise) {
if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind) {
have_int8_tensorwise = true;
}
if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) {
@@ -1051,7 +1051,7 @@ void StableDiffusionGGML::configure_weight_loading() {
} else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
if (have_int8_tensorwise) {
LOG_WARN(
"INT8 tensorwise weights do not support the immediately LoRA apply mode; "
"ComfyUI INT8/packed INT4 weights do not support the immediately LoRA apply mode; "
"using at_runtime instead");
apply_lora_immediately = false;
} else if (row_split_active()) {