diff --git a/docs/int4_convrot.md b/docs/int4_convrot.md new file mode 100644 index 000000000..05311bbc2 --- /dev/null +++ b/docs/int4_convrot.md @@ -0,0 +1,91 @@ +# INT4 Convrot Safetensors + +sd.cpp can load and execute ComfyUI packed int4 convrot safetensors directly. Two formats are supported: + +- `convrot_w4a4`: int4 weights with int4-class activations. +- `asym_w4a8_int8`: int4 codebook weights with int8 activations. + +The packed int4 weights are not converted to another weight type at load time. Weights consume roughly half a byte per parameter, and loading performs no reblocking step and writes no cache files. + +This requires the packed int4 convrot extensions in the patched GGML. +Builds with `SD_USE_UPSTREAM_GGML=ON` reject these files during loading. + +## Checkpoint format + +Each quantized linear module contains a packed weight tensor and a U8 quantization configuration blob: + +- `.weight`: a byte-packed I8 tensor holding two 4-bit values per byte, `K / 2` bytes per output row. +- `.comfy_quant`: a U8 tensor containing the JSON quantization configuration. + +A `convrot_w4a4` module additionally contains: + +- `.weight_scale`: one floating-point scale per output row. A two-dimensional `[out_features, 1]` representation is normalized to one dimension while loading. + +An `asym_w4a8_int8` module additionally contains: + +- `.weight_codebook`: 16 floating-point codebook entries shared by the module. +- `.weight_s_channel`: one floating-point scale per output row. +- `.weight_s_rel`: one F8_E4M3 relative scale per 16 input features per output row. + +The configuration blob for `convrot_w4a4` has this form: + +```json +{ + "format": "convrot_w4a4", + "convrot_groupsize": 64 +} +``` + +For `asym_w4a8_int8` the blob also carries the codebook quantization group size: + +```json +{ + "format": "asym_w4a8_int8", + "group_size": 16, + "convrot_groupsize": 64 +} +``` + +The convrot group size must be a power of four, must divide the input feature dimension, and must be `64` or `256`. The input feature dimension must additionally be a multiple of `32`. + +## How int4 convrot works + +The rotation follows the same scheme as INT8 convrot: a normalized regular Hadamard transform is applied offline to the weights and at runtime to the activations, preserving the linear operation while spreading outliers across each feature group. See [INT8 Convrot Safetensors](./int8_convrot.md) for the rotation details. + +For an original floating-point linear layer `Y = X W^T + b`, the checkpoint stores packed nibbles of the rotated weights. At runtime sd.cpp rotates and quantizes the activations once per input, reusing the packed result across all linear layers that share the same input and group size. + +### convrot_w4a4 + +The rotated weights are quantized to signed 4-bit values in `[-7, 7]` with one scale per output row. At runtime the activations are quantized to int8, which keeps the dot product on an exact integer grid: the 4-bit weights sign-extend to int8 without error, the accumulation is exact in 32-bit integers, and the output is reconstructed as + +```text +Y[r, o] ~= A[r, o] * s_x[r] * s_w[o] + b[o] +``` + +This variant introduces no additional weight error beyond the int4 quantization itself. + +### asym_w4a8_int8 + +Each nibble indexes a 16-entry per-module codebook. The effective weight is + +```text +W_rot[o, i] = codebook[code[o, i]] * s_channel[o] * s_rel[o, i / 16] +``` + +There is no integer grid for the codebook, so the kernels requantize the codebook to int8 when a block starts and compute the dot product with int8 integer instructions. The per-group relative scales are applied as a floating-point correction after each group. The result stays within a small tolerance of the floating-point reference instead of matching it bit for bit. + +## Backend support + +Both int4 formats run on CPU, NVIDIA CUDA, HIP (ROCm), and Vulkan, with convrot group sizes `64` and `256`. + +On HIP, plain `int8_tensorwise` models currently fall back to CPU execution; this is tracked by the separate INT8 HIP/BLAS changes and does not affect int4 models. + +LoRA adapters are applied at runtime without modifying the packed int4 weights. The int4 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing packed int4 weights. + +## Example + +ComfyUI int4 convrot safetensors can be passed to `--diffusion-model` without conversion: + +```powershell +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\flux-2-klein-4b_int4_convrot.safetensors --vae ..\models\vae\flux2-klein-vae.safetensors --llm ..\models\text_encoders\qwen_3_4b_int8_convrot.safetensors -p "a lovely cat holding a sign says 'sd.cpp'" --steps 8 --cfg-scale 1 --diffusion-fa -v --vae-tiling +``` diff --git a/ggml b/ggml index f583f393c..025d7a500 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit f583f393cd5dfdc129360bbf75cb3d49ccc837a8 +Subproject commit 025d7a5009ca59000652257cd3f992abdb3474e9 diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 74324169e..11bf6c10f 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -284,6 +284,41 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, #endif } +ggml_tensor* ggml_ext_linear_w4_convrot(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* weight_scales, + ggml_tensor* s_channel, + ggml_tensor* s_rel, + ggml_tensor* b, + int w4_convrot_kind, + int convrot_group_size, + float scale) { + GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f)); + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + + ggml_tensor* fused_bias = scale == 1.f ? b : nullptr; + if (x->ne[2] * x->ne[3] > 1024) { + int64_t ne2 = x->ne[2]; + int64_t ne3 = x->ne[3]; + x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); + x = ggml_mul_mat_w4_convrot(ctx, w, x, weight_scales, s_channel, s_rel, fused_bias, w4_convrot_kind, convrot_group_size); + x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); + } else { + x = ggml_mul_mat_w4_convrot(ctx, w, x, weight_scales, s_channel, s_rel, fused_bias, w4_convrot_kind, convrot_group_size); + } + + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + if (b != nullptr) { + x = ggml_add_inplace(ctx, x, b); + } + } + return x; +} + ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 33d8c3bf2..c06e2060a 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -80,6 +80,20 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, int convrot_group_size, float scale = 1.f); +// packed-w4 convrot linear; weight_scales is [N] F32 row scales for W4A4, or +// the [16] F32 codebook for W4A8 (with s_channel [N] and s_rel F8_E4M3 +// [K/16, N] supplied; both null for W4A4) +ggml_tensor* ggml_ext_linear_w4_convrot(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* weight_scales, + ggml_tensor* s_channel, + ggml_tensor* s_rel, + ggml_tensor* b, + int w4_convrot_kind, + int convrot_group_size, + float scale = 1.f); + ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 29a6808bd..25c6c1048 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -83,6 +83,7 @@ struct GGMLRunnerContext { std::function cache_tensor; std::function set_backend_tensor_data; std::map, ggml_tensor*> int8_convrot_cache; + std::map, ggml_tensor*> convrot_f32_cache; void capture_tensor(const std::string& name, ggml_tensor* tensor) { if (debug_tensors == nullptr || tensor == nullptr) { diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28..39ae7d85b 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -138,6 +138,41 @@ class Identity : public UnaryBlock { } }; +// Normalized regular Hadamard transform (H/gs^0.5) applied independently to each +// contiguous group of group_size elements along ne0, expressed as a butterfly of +// view/add/sub/concat stages so it runs on every backend without a dedicated op. +__STATIC_INLINE__ ggml_tensor* ggml_regular_hadamard_f32(ggml_context* ctx, ggml_tensor* x, int group_size) { + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(x)); + GGML_ASSERT(group_size > 1 && (group_size & (group_size - 1)) == 0); + GGML_ASSERT(x->ne[0] % group_size == 0); + const int64_t gs = group_size; + const int64_t n_rows = ggml_nelements(x) / x->ne[0]; + const int64_t m = x->ne[0] / gs; // independent Hadamard groups per row + ggml_tensor* y = ggml_reshape_3d(ctx, x, gs, m, n_rows); + for (int64_t b = gs; b > 1; b /= 2) { + const int64_t h = b / 2; + const int64_t blocks_per_group = (gs / b) * m; + if (ggml_nelements(y) != h * 2 * blocks_per_group * n_rows || !ggml_is_contiguous(y)) { + fprintf(stderr, "hadamard stage pre-check failed: b=%lld h=%lld y=[%lld,%lld,%lld,%lld] n_rows=%lld\n", + (long long)b, (long long)h, (long long)y->ne[0], (long long)y->ne[1], (long long)y->ne[2], + (long long)y->ne[3], (long long)n_rows); + GGML_ABORT("ggml_regular_hadamard_f32 stage shape mismatch"); + } + ggml_tensor* v = ggml_reshape_4d(ctx, y, h, 2, blocks_per_group, n_rows); + const size_t nb1 = ggml_row_size(GGML_TYPE_F32, h); + const size_t nb2 = ggml_row_size(GGML_TYPE_F32, 2 * h); + const size_t nb3 = ggml_row_size(GGML_TYPE_F32, 2 * h * blocks_per_group); + ggml_tensor* top = ggml_view_4d(ctx, v, h, 1, blocks_per_group, n_rows, nb1, nb2, nb3, 0); + ggml_tensor* bottom = ggml_view_4d(ctx, v, h, 1, blocks_per_group, n_rows, nb1, nb2, nb3, h * sizeof(float)); + ggml_tensor* sum = ggml_add(ctx, top, bottom); + ggml_tensor* diff = ggml_sub(ctx, top, bottom); + y = ggml_concat(ctx, sum, diff, 0); + } + y = ggml_scale(ctx, y, 1.0f / sqrtf((float)gs)); + return ggml_reshape_4d(ctx, y, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); +} + class Linear : public UnaryBlock { protected: int64_t in_features; @@ -148,6 +183,8 @@ class Linear : public UnaryBlock { bool has_weight_scale = false; bool int8_convrot = false; int int8_convrot_group_size = 0; + int w4_convrot_kind = W4_CONVROT_NONE; + int w4_convrot_group_size = 0; float scale; std::string prefix; @@ -156,19 +193,28 @@ class Linear : public UnaryBlock { has_weight_scale = false; int8_convrot = false; int8_convrot_group_size = 0; + w4_convrot_kind = W4_CONVROT_NONE; + w4_convrot_group_size = 0; + 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; + const bool is_w4_convrot = weight_storage != tensor_storage_map.end() && weight_storage->second.w4_convrot_kind != W4_CONVROT_NONE; enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); if (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 (is_w4_convrot) { + // packed nibbles: file tensor is [N, K/2] I8; the graph consumes the + // same bytes as [K/2, N] (ggml dims are reversed vs file layout) + GGML_ASSERT(in_features % 2 == 0); + params["weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_I8, in_features / 2, out_features); + } else { + params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features); + } if (bias) { - enum ggml_type wtype = GGML_TYPE_F32; - params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features); + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 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()) { + if (weight_scale_storage != tensor_storage_map.end() && !is_w4_convrot) { const int64_t scale_nelements = weight_scale_storage->second.nelements(); GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features); params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); @@ -180,6 +226,31 @@ class Linear : public UnaryBlock { int8_convrot = weight_storage->second.int8_convrot; int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; } + if (is_w4_convrot) { + w4_convrot_kind = weight_storage->second.w4_convrot_kind; + w4_convrot_group_size = weight_storage->second.w4_convrot_group_size; + if (w4_convrot_kind == W4_CONVROT_W4A4) { + // row scales load straight from the weight_scale companion + GGML_ASSERT(weight_scale_storage != tensor_storage_map.end()); + const int64_t scale_nelements = weight_scale_storage->second.nelements(); + GGML_ASSERT(scale_nelements == out_features); + params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); + has_weight_scale = true; + } else { + GGML_ASSERT(w4_convrot_kind == W4_CONVROT_W4A8); + GGML_ASSERT(in_features % 16 == 0); + auto codebook_storage = tensor_storage_map.find(prefix + "weight_codebook"); + auto s_channel_storage = tensor_storage_map.find(prefix + "weight_s_channel"); + auto s_rel_storage = tensor_storage_map.find(prefix + "weight_s_rel"); + GGML_ASSERT(codebook_storage != tensor_storage_map.end() && codebook_storage->second.nelements() == 16); + GGML_ASSERT(s_channel_storage != tensor_storage_map.end() && s_channel_storage->second.nelements() == out_features); + GGML_ASSERT(s_rel_storage != tensor_storage_map.end() && s_rel_storage->second.is_f8_e4m3 && + s_rel_storage->second.nelements() == out_features * (in_features / 16)); + 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); + params["weight_s_rel"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F8_E4M3, in_features / 16, out_features); + } + } } public: @@ -229,7 +300,7 @@ class Linear : public UnaryBlock { } ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; ggml_tensor* out = nullptr; - if (w->type == GGML_TYPE_I8) { + if (w->type == GGML_TYPE_I8 && w4_convrot_kind == W4_CONVROT_NONE) { if (x->type != GGML_TYPE_F32) { x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x); } @@ -274,6 +345,61 @@ class Linear : public UnaryBlock { } return out; } + if (w4_convrot_kind != W4_CONVROT_NONE) { + // packed-w4 kernel path: fused Hadamard+int8 activation quant, then + // the w4 mul_mat chain (mirrors the int8 branch above) + ggml_tensor* lora_input = x; + if (x->type != GGML_TYPE_F32) { + x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x); + } + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx->ggml_ctx, x); + } + if (ctx->weight_adapter && b != nullptr) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + // quantize only at unit scale: with scale != 1 the wrapper + // pre-scales the F32 input and the w4 kernels quantize inline + if (scale == 1.f) { + const auto cache_key = std::make_pair(x, w4_convrot_group_size); + auto cached = ctx->int8_convrot_cache.find(cache_key); + if (cached == ctx->int8_convrot_cache.end()) { + x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, w4_convrot_group_size); + ctx->int8_convrot_cache.emplace(cache_key, x); + } else { + x = cached->second; + } + } + ggml_tensor* weight_scales = nullptr; + ggml_tensor* s_channel = nullptr; + ggml_tensor* s_rel = nullptr; + if (w4_convrot_kind == W4_CONVROT_W4A4) { + weight_scales = params["weight_scale"]; + } else { + // the kernel folds codebook x s_channel into its per-group LUT + weight_scales = params["weight_codebook"]; + s_channel = params["weight_s_channel"]; + s_rel = params["weight_s_rel"]; + } + out = ggml_ext_linear_w4_convrot(ctx->ggml_ctx, x, w, weight_scales, s_channel, s_rel, b, + w4_convrot_kind, w4_convrot_group_size, scale); + if (ctx->weight_adapter) { + // LoRA was trained on unrotated activations; apply it separately + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; + forward_params.linear.force_prec_f32 = force_prec_f32; + forward_params.linear.scale = scale; + out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, + ctx->backend, + lora_input, + w, + out, + prefix, + forward_params); + return out; + } + return out; + } if (has_weight_scale) { out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale); out = ggml_mul(ctx->ggml_ctx, out, weight_scale); diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 3e6974237..0f3dc25ca 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; // convrot_groupsize: Hadamard rotation block size + int w4a8_group_size = 0; // asym_w4a8_int8 group_size: s_rel group axis (16) }; static bool read_comfy_quant_config(std::ifstream& file, @@ -140,6 +141,7 @@ 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.w4a8_group_size = json.value("group_size", 0); } catch (const std::exception&) { set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'"); return false; @@ -340,29 +342,64 @@ 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 (type != GGML_TYPE_I8) { - set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'"); - return false; - } - 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 (config != comfy_quant_configs.end()) { + const std::string& format = config->second.format; + const bool is_w4 = format == "convrot_w4a4" || format == "asym_w4a8_int8"; + if (config->second.format == "int8_tensorwise") { + if (type != GGML_TYPE_I8) { + set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'"); + return false; + } + 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) { + set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'"); + return false; + } + } + tensor_storage.is_int8_tensorwise = true; + tensor_storage.int8_convrot = config->second.convrot; + tensor_storage.int8_convrot_group_size = config->second.group_size; + } else if (is_w4) { + // packed I8 [N, K/2]; ne is ggml order here: ne[0] = K/2, ne[1] = N + if (type != GGML_TYPE_I8) { + set_error(error, "ComfyUI " + format + " weight is not I8: '" + name + "'"); + return false; + } + if (tensor_storage.n_dims != 2) { + set_error(error, "ComfyUI " + format + " weight is not 2D: '" + name + "'"); + return false; } - if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) { + const int64_t packed_k = tensor_storage.ne[0]; + const int convrot_gs = config->second.group_size > 0 ? config->second.group_size : 256; + int gs_remainder = convrot_gs; + while (gs_remainder > 1 && gs_remainder % 4 == 0) { + gs_remainder /= 4; + } + // K must fill whole 32-blocks (converter contract), and K % GS == 0 + if (gs_remainder != 1 || (2 * packed_k) % 32 != 0 || (2 * packed_k) % convrot_gs != 0) { set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'"); return false; } + tensor_storage.w4_convrot_kind = (format == "convrot_w4a4") ? W4_CONVROT_W4A4 : W4_CONVROT_W4A8; + tensor_storage.w4_convrot_group_size = convrot_gs; + if (format == "asym_w4a8_int8") { + if (config->second.w4a8_group_size <= 0 || (2 * packed_k) % config->second.w4a8_group_size != 0) { + set_error(error, "invalid asym_w4a8_int8 group size for tensor '" + name + "'"); + return false; + } + tensor_storage.w4a8_group_size = config->second.w4a8_group_size; + } } - tensor_storage.is_int8_tensorwise = true; - 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()); 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") && tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) { tensor_storage.ne[0] = tensor_storage.ne[1]; tensor_storage.ne[1] = 1; @@ -405,6 +442,8 @@ bool read_safetensors_file(const std::string& file_path, return false; } + // Packed w4 convrot weights stay raw: the w4 kernels consume the packed + // bytes in place (type I8, ne = [K/2, N] in ggml order). No rewrite. tensor_storages.push_back(tensor_storage); // LOG_VERBOSE("%s %s", tensor_storage.to_string().c_str(), dtype.c_str()); diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 1916d1f5e..285a02d56 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -13,6 +13,15 @@ #define SD_MAX_DIMS 5 +// ComfyUI packed 4-bit convrot weight formats. The storage is rewritten at read +// time to its converted representation (de-packed ne + target ggml type), the +// reblocked bytes are produced by ModelLoader before tensor loading. +enum W4ConvRotKind { + W4_CONVROT_NONE = 0, + W4_CONVROT_W4A4 = 1, // convrot_w4a4: signed nibbles + per-row F32 scale -> GGML_TYPE_Q4_0 + W4_CONVROT_W4A8 = 2, // asym_w4a8_int8: codebook + s_channel + s_rel (F8) -> GGML_TYPE_IQ4_NL +}; + struct TensorStorage { std::string name; ggml_type type = GGML_TYPE_F32; @@ -24,6 +33,9 @@ struct TensorStorage { bool is_int8_tensorwise = false; bool int8_convrot = false; int int8_convrot_group_size = 0; + int w4_convrot_kind = W4_CONVROT_NONE; + int w4_convrot_group_size = 0; + int w4a8_group_size = 0; int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; int n_dims = 0; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ee5aa49e3..57e1dc7ab 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -791,6 +791,10 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru tensor_type_rules_ = tensor_type_rules; auto map_rules = parse_tensor_type_rules(tensor_type_rules); for (auto& [name, tensor_storage] : tensor_storage_map) { + if (tensor_storage.w4_convrot_kind != W4_CONVROT_NONE) { + // target type is intrinsic to the packed format, not subject to overrides + continue; + } tensor_storage.expected_type = GGML_TYPE_COUNT; ggml_type dst_type = wtype; for (const auto& tensor_type_rule : map_rules) { @@ -1513,7 +1517,15 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage if (tensor_storage.is_int8_tensorwise) { return false; } - if (type != GGML_TYPE_COUNT) { + // Packed ComfyUI w4 convrot weights are consumed raw by the w4 kernels; + // never route them through the generic conversion machinery. + if (tensor_storage.w4_convrot_kind != W4_CONVROT_NONE) { + return false; + } + if (ends_with(name, ".weight_codebook") || ends_with(name, ".weight_s_rel") || + ends_with(name, ".weight_s_channel")) { + // Pass, do not convert (w4 convrot companions) + } else if (type != GGML_TYPE_COUNT) { if (ggml_is_quantized(type) && tensor_storage.ne[0] % ggml_blck_size(type) != 0) { // Pass, do not convert } else if (ends_with(name, ".bias")) { diff --git a/src/model_loader.h b/src/model_loader.h index c5335b53a..1d42c4336 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -144,6 +144,9 @@ class ModelLoader { bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const; int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT) const; + +private: +public: ~ModelLoader() = default; }; diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index b5dfe7c7c..6a793044a 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -12,7 +12,9 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) { 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.int8_convrot_group_size == b.int8_convrot_group_size; + a.int8_convrot_group_size == b.int8_convrot_group_size && + a.w4_convrot_kind == b.w4_convrot_kind && a.w4_convrot_group_size == b.w4_convrot_group_size && + a.w4a8_group_size == b.w4a8_group_size; } void ModelManager::invalidate_sources(const std::unordered_set& states) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e41c1151a..46187e3f0 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1010,7 +1010,7 @@ void StableDiffusionGGML::configure_weight_loading() { const auto wtype_stat = model_loader.get_wtype_stat(); bool have_int8_tensorwise = 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 != W4_CONVROT_NONE) { have_int8_tensorwise = true; break; }