diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 5c059892e8..77075989dd 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -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 | diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index e8036c60ea..49e77bcefc 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -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 diff --git a/docs/model_config.md b/docs/model_config.md index 0d9290c1d5..d42c3226ce 100644 --- a/docs/model_config.md +++ b/docs/model_config.md @@ -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 diff --git a/src/core/cpu/registration.cpp b/src/core/cpu/registration.cpp index d63fd1f099..16e13dfe22 100644 --- a/src/core/cpu/registration.cpp +++ b/src/core/cpu/registration.cpp @@ -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); } diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h index eba629c2a6..77f4bacfb0 100644 --- a/src/core/cpu/registration.h +++ b/src/core/cpu/registration.h @@ -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 diff --git a/src/core/cpu/w4.cpp b/src/core/cpu/w4.cpp new file mode 100644 index 0000000000..4da0f2b51d --- /dev/null +++ b/src/core/cpu/w4.cpp @@ -0,0 +1,95 @@ +#include "core/cpu/registration.h" +#include "core/ggml_extensions.h" + +#ifndef SD_USE_UPSTREAM_GGML +#include + +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(data); + const auto& op = e.op; + const int64_t k = op.input->ne[0], n = op.weight->ne[1]; + const auto* w = static_cast(op.weight->data); + const auto* x = static_cast(op.input->data); + const auto* xs = reinterpret_cast(x + k * ((e.rows + 3) / 4 * 4)); + const auto* ws = static_cast(op.weight_scale->data); + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const auto* sc = op.s_channel ? static_cast(op.s_channel->data) : nullptr; + const auto* sr = op.s_rel ? static_cast(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_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(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 diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 0a95d6cf7c..593445bf0e 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -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() diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 45f6ecdf14..60f36bb1e9 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -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 + static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) { + const int64_t group = blockIdx.x, groups = k / 64, row = group / groups; + const int tid = threadIdx.x; + __shared__ float values[64]; + __shared__ float maxima[16]; +#pragma unroll + for (int j = 0; j < 4; ++j) + values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f); + __syncthreads(); +#pragma unroll + for (int stride = 1; stride < 64; stride *= 4) { + const int i = (tid / stride) * 4 * stride + tid % stride; + const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride]; + values[i] = convrot_h4_cuda(a, b, c, d, 0); + values[i + stride] = convrot_h4_cuda(a, b, c, d, 1); + values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2); + values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3); + __syncthreads(); + } + if constexpr (Quantize) { + const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int value = __float2int_rn(values[tid + j * 16] * inverse); + dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value))); + } + } else { + float amax = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) + amax = fmaxf(amax, fabsf(values[tid + j * 16])); + maxima[tid] = amax; + __syncthreads(); + for (int stride = 8; stride; stride >>= 1) { + if (tid < stride) + maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); + __syncthreads(); + } + if (!tid) + partial[group] = maxima[0]; + } + } static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { int group; - if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0])) + if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0])) return false; const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]); - return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0; + return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0; } enum Kind { MAT, QUANTIZE, @@ -421,7 +464,8 @@ namespace SD_CUDA_NAMESPACE { return 0; if (first + 2 < ggml_graph_n_nodes(g)) { const auto* quant = ggml_graph_node(g, first + 2); - if (quant->src[0] == next && convrot_supports(device, quant, nullptr) && + int group; + if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) && size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) return 3; } @@ -430,6 +474,7 @@ namespace SD_CUDA_NAMESPACE { struct Plan { ggml_ext_gpu_device_info device; Kind kind; + int group; int64_t k, n, rows, padded; size_t quant, scales, accum, partial, output, gemm_offset; SDI8CudaGemmPlan* gemm = nullptr; @@ -456,7 +501,10 @@ namespace SD_CUDA_NAMESPACE { p->n = mat ? nodes[0]->ne[0] : 0; p->rows = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]); p->padded = (p->rows + 3) / 4 * 4; - size_t total = 0; + p->group = mat ? op.group_size : 0; + if (!mat) + sd_ggml_convrot_describe(dst, &p->group); + size_t total = 0; if (mat) { p->quant = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0); p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0); @@ -480,7 +528,8 @@ namespace SD_CUDA_NAMESPACE { p->accum = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0); } const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory; - p->partial = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); + p->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) + : 0); p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); *workspace = total; return GGML_STATUS_SUCCESS; @@ -506,6 +555,17 @@ namespace SD_CUDA_NAMESPACE { return error; } const size_t shared = size_t(p.k) * sizeof(float); + if (p.group == 1) { + quantize_rowwise_i8_cuda<<>>(src, dst, scales, p.k, p.rows); + return cuda_last_error(); + } + if (p.group == 64) { + const auto groups = p.k / 64; + convrot_h64<<>>(src, dst, partial, scales, p.k); + reduce_convrot_row_amax_i8_cuda<<>>(partial, scales, groups, p.rows); + convrot_h64<<>>(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 configured; diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index d75728021f..1f48174313 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -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); diff --git a/src/core/cuda/w4.cu b/src/core/cuda/w4.cu new file mode 100644 index 0000000000..ef46d02369 --- /dev/null +++ b/src/core/cuda/w4.cu @@ -0,0 +1,873 @@ +#include +#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(&a); + const auto* bv = reinterpret_cast(&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(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(dst->data), b = reinterpret_cast(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(dst->data), b = reinterpret_cast(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(context.stream); + const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0]; + auto* scratch = e->workspace ? static_cast(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(dst->data); + const auto* w = static_cast(op.weight->data); + const auto* x = static_cast(op.input->data); + const auto* xs = reinterpret_cast(x + k * ((rows + 3) / 4 * 4)); + const auto* ws = static_cast(op.weight_scale->data); + const auto* bias = op.bias ? static_cast(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(scratch + (alias ? ggml_nelements(dst) : 0)); + mul_mat_w4a4_convrot_dp4a<<>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2)); + if (cuda_last_error() != cuda_success) + return GGML_STATUS_FAILED; + dequantize<<>>(accum, xs, ws, bias, output, n, n * rows); + } else { + mul_mat_w4a8_convrot<<>>(w, x, ws, + static_cast(op.s_channel->data), static_cast(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)); + } +} diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 27089f6c3b..25fa22265f 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -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; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index 86c9121a7a..bc9ee9a5c4 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -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, diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index 55acf0bb91..e71bcd5339 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -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) diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index 7e3e67d48d..5ac8151854 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -11,7 +11,10 @@ #include "mul_mat_i8_tensorwise.h" #include "mul_mat_i8_tensorwise_cm1.h" #include "mul_mat_i8_tensorwise_cm2.h" +#include "mul_mat_w4a4_convrot.h" +#include "mul_mat_w4a8_convrot.h" #include "quantize_i8_convrot.h" +#include "quantize_i8_convrot_h64.h" #include "quantize_i8_convrot_register.h" #include "quantize_i8_convrot_subgroup.h" #include "rms_modulate_convrot.h" @@ -111,6 +114,23 @@ namespace sd_vulkan_i8 { (cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] && d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16; } + static bool w4_supports(ggml_backend_dev_t dev, const ggml_tensor* dst, void*) { + Device d; + SDW4Operation op; + if (!device_info(dev, d) || !sd_ggml_w4_describe(dst, &op) || !bounded(d, dst) || + sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot) <= 4 || sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot) <= 4 || + d.compute.shared_memory < (op.kind == SDW4Kind::W4A8 ? 16452u : 13312u) || d.info.max_push_constants < 16) + return false; + for (auto* t : dst->src) + if (t && (!bounded(d, t) || ggml_nbytes(t) > UINT32_MAX)) + return false; + const uint64_t k = op.input->ne[0], n = dst->ne[0], rows = ggml_nrows(dst); + return k % 64 == 0 && k <= UINT32_MAX && n <= UINT32_MAX && rows <= UINT32_MAX && + ggml_nelements(dst) <= UINT32_MAX && + (k * ((rows + 3) / 4 * 4)) % d.info.buffer_alignment == 0 && + (n + 63) / 64 <= d.info.max_group_count[0] && (rows + 31) / 32 <= d.info.max_group_count[1] && + d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16; + } static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) { if (i + 2 >= count_nodes(graph) || !device.subgroup) { return false; @@ -534,9 +554,12 @@ namespace sd_vulkan_i8 { status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups); } else { ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}}; - auto pipeline = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters); - bool subgroup = d.subgroup && k >= 1024; - uint32_t width = k; + int group; + if (!sd_ggml_convrot_describe(output, &group)) + return GGML_STATUS_FAILED; + auto pipeline = group == 64 ? SD_VK_PIPELINE(quantize_i8_convrot_h64, 3, QuantParameters) : SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters); + bool subgroup = group == 256 && d.subgroup && k >= 1024; + uint32_t width = k; if (subgroup) { pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters); if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) { @@ -653,6 +676,42 @@ namespace sd_vulkan_i8 { } return finish(e, d, output, dst, alias, status); } + static ggml_status multiply_w4(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* output) { + SDW4Operation op; + if (!sd_ggml_w4_describe(output, &op)) + return GGML_STATUS_FAILED; + const uint32_t k = uint32_t(op.input->ne[0]), n = uint32_t(output->ne[0]), rows = uint32_t(ggml_nrows(output)); + const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4); + const ggml_ext_vulkan_binding xs = {op.input, nullptr, scale_offset, size_t(rows) * 4}; + std::vector 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(); diff --git a/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp b/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp new file mode 100644 index 0000000000..f5fb10e6f4 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp @@ -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; + } + } + } + } +} diff --git a/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp b/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp new file mode 100644 index 0000000000..4438ce2876 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp @@ -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; + } + } + } + } +} diff --git a/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp b/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp new file mode 100644 index 0000000000..2fd6dbb463 --- /dev/null +++ b/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp @@ -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); + } +} diff --git a/src/model/adapter/ip_adapter.hpp b/src/model/adapter/ip_adapter.hpp index 9f8dc9f64f..b36977df4e 100644 --- a/src/model/adapter/ip_adapter.hpp +++ b/src/model/adapter/ip_adapter.hpp @@ -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; diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index b2657a3691..7b9e8ab82d 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -1127,6 +1127,16 @@ struct MultiLoraAdapter : public WeightAdapter { 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) { diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index e79ee2f12a..66180fa8d1 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -319,7 +319,7 @@ class CrossAttention : public GGMLBlock { 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(new Linear(ip_dim, inner_dim, false)); blocks["to_v_ip"] = std::shared_ptr(new Linear(ip_dim, inner_dim, false)); } diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 7529848477..0653ae1fff 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -148,6 +148,7 @@ class Linear : public UnaryBlock { 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 @@ class Linear : public UnaryBlock { 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 @@ class Linear : public UnaryBlock { 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; } @@ -259,13 +272,31 @@ class Linear : public UnaryBlock { } } #endif - out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, - x, - w, - weight_scale, - b, - int8_convrot ? int8_convrot_group_size : 0, - scale); +#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(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, + w, + weight_scale, + b, + int8_convrot ? int8_convrot_group_size : 0, + scale); if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 3f652e1589..ecdb7def3e 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -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]; diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 15badb9834..5898ac0679 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -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) { diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index 0c42235f8f..2a65ea1d9c 100644 --- a/src/model/diffusion/fibo2.hpp +++ b/src/model/diffusion/fibo2.hpp @@ -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]; diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 40edca0ed9..843e70b458 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -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")) { diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index fac1b1a567..105b22565d 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -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; } diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index fba817aa67..e44c5a9c6e 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -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; diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index bcfefd97bf..0f8c816739 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -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(tensor_storage.ne[0] / config.joint_attention_dim); + config.selected_layer_count = static_cast(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; diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index d15d754080..b94d7fdb86 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -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) { diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index 4a9ea1e5ab..6b2b534732 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -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) { diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index 32e38012bd..f5d2b36157 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -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; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index a7c2f1f930..b5aec4397d 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -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")) { diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index c6e90dae1c..5029759d33 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -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]; diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 613cfc11a7..b032273db5 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -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(config.patch_size) * config.patch_size; if (patch_area > 0) { config.out_channels = tensor_storage.ne[1] / patch_area; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index a0fae7448f..2641417abf 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -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(depth, std::stoi(key.substr(7)) + 1); diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp index dc6e2c2e7d..74bde6d453 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -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]; diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 76cbccb838..9452ec4678 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -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]; diff --git a/src/model/diffusion/sefi_image.hpp b/src/model/diffusion/sefi_image.hpp index 6782b68b49..7781532260 100644 --- a/src/model/diffusion/sefi_image.hpp +++ b/src/model/diffusion/sefi_image.hpp @@ -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; } } diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 7c0ff09506..675d94296c 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -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(time_embed->ne[0]); + config.model_channels = static_cast(time_embed->logical_ne0()); config.time_embed_dim = static_cast(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(label_emb->ne[0]); + config.adm_in_channels = static_cast(label_emb->logical_ne0()); config.time_embed_dim = static_cast(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(tensor_storage.ne[0]); + config.context_dim = static_cast(tensor_storage.logical_ne0()); break; } } diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index 1c66c96b7f..bc9ddc4498 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -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]; diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp index f44f2ff9f8..2ba6680329 100644 --- a/src/model/diffusion/z_image_l2p.hpp +++ b/src/model/diffusion/z_image_l2p.hpp @@ -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(std::lround(std::sqrt(static_cast(x_embedder->second.ne[0] / config.in_channels)))); + config.patch_size = static_cast(std::lround(std::sqrt(static_cast(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. diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 7e7ac30984..29c6f87311 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -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]; } } diff --git a/src/model/upscaler/h3_latent_upscaler.hpp b/src/model/upscaler/h3_latent_upscaler.hpp index 006415081f..80b96e7395 100644 --- a/src/model/upscaler/h3_latent_upscaler.hpp +++ b/src/model/upscaler/h3_latent_upscaler.hpp @@ -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(embed_mid->ne[0]); + config.embed_dim = static_cast(embed_mid->logical_ne0()); } std::set in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx; diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index cf424af921..058606267d 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -109,8 +109,9 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { struct ComfyQuantConfig { std::string format; - bool convrot = false; - int group_size = 0; + 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& 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; diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 1916d1f5e5..1062174ebc 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -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); } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 407885bbff..73025a4075 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -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& 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) { diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index 10c7d2e68f..adeea95d22 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -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; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 98e80e1174..4875858b82 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -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()) {