diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2d83c0d4a..884ce05bb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -369,7 +369,7 @@ jobs: matrix: include: - build: "cpu" - defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" + defines: "-DGGML_NATIVE=OFF -DSD_BUILD_TESTS=ON -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "cuda12" defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "vulkan" @@ -421,6 +421,10 @@ jobs: cmake .. -DCMAKE_CXX_FLAGS='/bigobj' -G Ninja -DCMAKE_C_COMPILER=cl.exe -DCMAKE_CXX_COMPILER=cl.exe -DCMAKE_BUILD_TYPE=Release ${{ matrix.defines }} cmake --build . + - name: Test ConvRot safetensors loader + if: ${{ matrix.build == 'cpu' }} + run: ctest --test-dir build -R "^test-safetensors-convrot$" --output-on-failure + - name: Get commit hash id: commit if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} diff --git a/ggml b/ggml index 7d9ce11cd..552c92298 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 7d9ce11cd47f338b361a00e866ffe7c224abedff +Subproject commit 552c92298881f2b8fc99d938abc3a426b90f3497 diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index b29803d92..d2e783d22 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -5,8 +5,8 @@ #include #include #include -#include #include +#include #include #include #include @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -43,6 +44,81 @@ #define EPS 1e-05f +// Construct only the operation metadata needed for the normal backend +// supports_op query. This stays private to the loader policy: backend +// capabilities are expressed through the existing ggml interface, not a new +// public ConvRot-specific API. +inline bool ggml_backend_supports_convrot_op(ggml_backend_t backend) { + if (backend == nullptr) { + return false; + } + std::vector storage(4 * ggml_tensor_overhead() + 1024); + ggml_init_params params = { + /*.mem_size =*/ storage.size(), + /*.mem_buffer =*/ storage.data(), + /*.no_alloc =*/ true, + }; + ggml_context* ctx = ggml_init(params); + if (ctx == nullptr) { + return false; + } + ggml_tensor* activations = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); + ggml_tensor* weights = ggml_new_tensor_2d(ctx, GGML_TYPE_I8, 256, 1); + ggml_tensor* scales = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_tensor* op = ggml_mul_mat_convrot(ctx, activations, weights, scales, 256); + const bool supported = ggml_backend_supports_op(backend, op); + ggml_free(ctx); + return supported; +} + +// Select the compact representation before any model parameter tensor is +// created. The default is deliberately native: an unsupported backend is a +// configuration error rather than a silent CPU reroute or full F16 expansion. +// Set SD_CONVROT_MODE=compat to explicitly request the compatibility loader. +inline String2TensorStorage select_convrot_tensor_storage(ggml_backend_t backend, + const String2TensorStorage& source, + const std::string& component, + const std::string& prefix = "") { + // OrderedMap's default copy also copies its iterator index; rebuild it so + // this independent policy view owns a valid index into its own list. + String2TensorStorage selected; + bool has_convrot = false; + for (const auto& [name, storage] : source) { + selected.insert({name, storage}); + has_convrot = has_convrot || + (name.rfind(prefix, 0) == 0 && storage.is_comfy_int8_convrot_weight()); + } + if (!has_convrot) { + return selected; + } + + const char* mode = std::getenv("SD_CONVROT_MODE"); + const bool compatibility_mode = mode != nullptr && std::strcmp(mode, "compat") == 0; + if (mode != nullptr && !compatibility_mode && std::strcmp(mode, "native") != 0) { + throw std::runtime_error("invalid SD_CONVROT_MODE; expected 'native' or 'compat'"); + } + const char* backend_name = backend != nullptr ? ggml_backend_name(backend) : "unknown"; + if (compatibility_mode) { + LOG_INFO("ConvRot: using explicitly selected F16 compatibility path for %s on backend %s", + component.c_str(), backend_name); + return selected; + } + if (!ggml_backend_supports_convrot_op(backend)) { + throw std::runtime_error("ConvRot native support is required for " + component + + " but backend '" + backend_name + + "' lacks the 256-wide I8/F32 ConvRot operation; use a capable backend or set " + "SD_CONVROT_MODE=compat to select the F16 compatibility path"); + } + for (auto& [name, storage] : selected) { + if (name.rfind(prefix, 0) == 0 && storage.is_comfy_int8_convrot_weight()) { + storage.comfy_int8_native_enabled = true; + } + } + LOG_INFO("ConvRot: selected native compact I8/F32 path for %s on backend %s", + component.c_str(), backend_name); + return selected; +} + #ifndef __STATIC_INLINE__ #define __STATIC_INLINE__ static inline #endif @@ -1695,6 +1771,15 @@ struct WeightAdapter { ggml_tensor* b, const std::string& prefix, ForwardParams forward_params) = 0; + // Return only the adapter's output-space contribution. Native operations + // such as compact ConvRot own their base-weight arithmetic and therefore + // cannot use forward_with_lora() without recomputing an incompatible base. + virtual ggml_tensor* lora_output_delta(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + const std::string& prefix, + ForwardParams forward_params) = 0; virtual size_t get_extra_graph_size() = 0; }; @@ -3875,12 +3960,31 @@ class Linear : public UnaryBlock { bool force_prec_f32; bool allow_weight_scale; bool has_weight_scale = false; + // This is distinct from `weight_scale`: the latter is a regular + // post-linear model parameter, while ConvRot's F32 vector is a private + // sidecar input to GGML_OP_MUL_MAT_CONVROT. + bool has_convrot_weight = false; + bool use_convrot_f16_compat = false; float scale; std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { this->prefix = prefix; has_weight_scale = false; + has_convrot_weight = false; + use_convrot_f16_compat = false; + const auto storage_it = tensor_storage_map.find(prefix + "weight"); + if (storage_it != tensor_storage_map.end() && storage_it->second.is_comfy_int8_convrot_weight() && + storage_it->second.comfy_int8_native_enabled) { + params["weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_I8, in_features, out_features); + params["weight.convrot_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features); + has_convrot_weight = true; + use_convrot_f16_compat = storage_it->second.name.rfind("text_encoders.llm.", 0) == 0; + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features); + } + return; + } 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; @@ -3928,7 +4032,33 @@ class Linear : public UnaryBlock { } ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; ggml_tensor* out = nullptr; - if (ctx->weight_adapter) { + if (has_convrot_weight) { + // ConvRot weights and their tensor-wise scales remain compact at + // rest. The operator owns the scale semantics; do not route it + // through the ordinary `weight_scale` post-multiply path. + out = ggml_mul_mat_convrot(ctx->ggml_ctx, x, w, params["weight.convrot_scale"], 256); + // MiniMax H3's ConvRot text encoder is calibrated for the F16 + // compatibility arithmetic. CUDA reconstructs one F16 matrix at + // a time and uses its standard F16 GEMM without retaining an F16 + // copy of the complete text encoder. Other backends may ignore + // this hint and keep their native compact implementation. + if (use_convrot_f16_compat) { + ggml_mul_mat_convrot_set_f16_compat(out, true); + } + if (b != nullptr) { + out = ggml_add_inplace(ctx->ggml_ctx, out, b); + } + if (ctx->weight_adapter) { + 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; + if (ggml_tensor* delta = ctx->weight_adapter->lora_output_delta( + ctx->ggml_ctx, ctx->backend, x, w, prefix, forward_params)) { + out = ggml_add_inplace(ctx->ggml_ctx, out, delta); + } + } + } else if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 6c0f41a0a..c076b189d 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -1300,6 +1300,22 @@ struct MultiLoraAdapter : public WeightAdapter { return out; } + ggml_tensor* lora_output_delta(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + const std::string& prefix, + WeightAdapter::ForwardParams forward_params) override { + ggml_tensor* delta = nullptr; + for (auto& lora_model : lora_models) { + ggml_tensor* current = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight"); + if (current != nullptr) { + delta = delta == nullptr ? current : ggml_add_inplace(ctx, delta, current); + } + } + return delta; + } + size_t get_extra_graph_size() override { size_t lora_tensor_num = 0; for (auto& lora_model : lora_models) { diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index e5a5a9a9b..a372fa9d2 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -162,10 +162,7 @@ namespace MiniMaxH3 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); - auto uv = ggml_ext_chunk(ctx->ggml_ctx, fc1->forward(ctx, x), 2, 0); - return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx, - ggml_silu(ctx->ggml_ctx, uv[0]), - uv[1])); + return fc2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, fc1->forward(ctx, x))); } }; @@ -176,13 +173,30 @@ namespace MiniMaxH3 { static ggml_tensor* apply_partial_rope(ggml_context* ctx, ggml_tensor* x, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_backend_t backend = nullptr) { int64_t rot_dim = pe->ne[2] * 2; GGML_ASSERT(rot_dim <= x->ne[0]); - auto rotated = Rope::apply_rope(ctx, - ggml_ext_slice(ctx, x, 0, 0, rot_dim), - pe, - false); + auto input = ggml_ext_slice(ctx, x, 0, 0, rot_dim); + ggml_tensor* rotated = nullptr; + static const bool fused_disabled = std::getenv("GGML_ROPE_FLUX_DISABLE") != nullptr; + if (backend != nullptr && !fused_disabled) { + // H3 pairs the first and second halves of the rotary dimensions; + // ROPE_FLUX pairs adjacent values. Pack/unpack that ordering around + // the fused operation, without expanding PE across every head. + auto paired = ggml_reshape_4d(ctx, input, rot_dim / 2, 2, x->ne[1], x->ne[2] * x->ne[3]); + paired = ggml_cont(ctx, ggml_permute(ctx, paired, 1, 0, 2, 3)); + paired = ggml_reshape_4d(ctx, paired, rot_dim, x->ne[1], x->ne[2], x->ne[3]); + auto fused = ggml_rope_flux(ctx, paired, pe); + if (ggml_backend_supports_op(backend, fused)) { + rotated = ggml_reshape_4d(ctx, fused, 2, rot_dim / 2, x->ne[2], x->ne[1] * x->ne[3]); + rotated = ggml_cont(ctx, ggml_permute(ctx, rotated, 1, 0, 2, 3)); + rotated = ggml_reshape_3d(ctx, rotated, rot_dim, x->ne[2], x->ne[1] * x->ne[3]); + } + } + if (rotated == nullptr) { + rotated = Rope::apply_rope(ctx, input, pe, false); + } if (rot_dim == x->ne[0]) { return rotated; } @@ -223,8 +237,8 @@ namespace MiniMaxH3 { q = q_norm->forward(ctx, q); k = k_norm->forward(ctx, k); if (pe != nullptr) { - q = apply_partial_rope(ctx->ggml_ctx, q, pe); - k = apply_partial_rope(ctx->ggml_ctx, k, pe); + q = apply_partial_rope(ctx->ggml_ctx, q, pe, ctx->backend); + k = apply_partial_rope(ctx->ggml_ctx, k, pe, ctx->backend); } else { q = attention_layout(ctx->ggml_ctx, q); k = attention_layout(ctx->ggml_ctx, k); @@ -986,7 +1000,9 @@ namespace MiniMaxH3 { : DiffusionModelRunner(backend, prefix, weight_manager), config(Config::detect_from_weights(tensors, prefix)), model(config) { - model.init(params_ctx, tensors, prefix); + model.init(params_ctx, + select_convrot_tensor_storage(backend, tensors, "MiniMax-H3 diffusion model", prefix), + prefix); } std::string get_desc() override { diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index f4dfa9f76..0615546b3 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -1771,7 +1771,9 @@ namespace LLM { } } model = LLM(config, enable_vision, config.llama_cpp_style); - model.init(params_ctx, tensor_storage_map, prefix); + model.init(params_ctx, + select_convrot_tensor_storage(backend, tensor_storage_map, "LLM", prefix), + prefix); } std::string get_desc() override { diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index df71eab11..c8392e078 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -5,7 +5,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -93,10 +96,200 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { ttype = GGML_TYPE_I32; } else if (dtype == "I64") { ttype = GGML_TYPE_I32; + } else if (dtype == "I8") { + ttype = GGML_TYPE_I8; } return ttype; } +struct SafetensorsTensorInfo { + std::string dtype; + std::vector shape; + uint64_t begin = 0; + uint64_t end = 0; +}; + +struct ComfyInt8Info { + bool convrot = false; + uint32_t group_size = 0; + TensorStorageSidecar scale; +}; + +static bool read_safetensors_tensor_info(const nlohmann::json& value, + const std::string& name, + uint64_t data_size, + SafetensorsTensorInfo* result, + std::string* error) { + try { + if (!value.is_object() || !value.contains("dtype") || !value["dtype"].is_string() || + !value.contains("shape") || !value["shape"].is_array() || + !value.contains("data_offsets") || !value["data_offsets"].is_array() || + value["data_offsets"].size() != 2) { + set_error(error, "invalid safetensors descriptor for tensor '" + name + "'"); + return false; + } + + SafetensorsTensorInfo info; + info.dtype = value["dtype"].get(); + if (value["shape"].size() > SD_MAX_DIMS) { + set_error(error, "too many dimensions for tensor '" + name + "'"); + return false; + } + uint64_t elements = 1; + for (const auto& dimension : value["shape"]) { + int64_t size = dimension.get(); + if (size <= 0 || elements > static_cast(std::numeric_limits::max()) / + static_cast(size)) { + set_error(error, "invalid dimension for tensor '" + name + "'"); + return false; + } + elements *= static_cast(size); + info.shape.push_back(size); + } + info.begin = value["data_offsets"][0].get(); + info.end = value["data_offsets"][1].get(); + if (info.begin > info.end || info.end > data_size) { + set_error(error, "data offsets out of bounds for tensor '" + name + "'"); + return false; + } + *result = std::move(info); + return true; + } catch (const std::exception&) { + set_error(error, "invalid safetensors descriptor for tensor '" + name + "'"); + return false; + } +} + +static bool is_power_of_four(uint64_t value) { + if (value < 4) { + return false; + } + while (value % 4 == 0) { + value /= 4; + } + return value == 1; +} + +static bool read_comfy_int8_metadata(std::ifstream& file, + const std::map& tensors, + uint64_t data_start, + std::map* result, + std::set* scale_tensor_names, + std::string* error) { + result->clear(); + scale_tensor_names->clear(); + constexpr const char* marker_suffix = ".comfy_quant"; + constexpr size_t marker_suffix_len = 12; + + for (const auto& [marker_name, marker] : tensors) { + if (!ends_with(marker_name, marker_suffix)) { + continue; + } + if (marker.dtype != "U8") { + set_error(error, "ComfyUI quantization marker '" + marker_name + "' must use U8 storage"); + return false; + } + if (marker.end - marker.begin > 4096) { + set_error(error, "ComfyUI quantization marker '" + marker_name + "' is too large"); + return false; + } + uint64_t marker_elements = 1; + for (int64_t dimension : marker.shape) { + if (marker_elements > std::numeric_limits::max() / + static_cast(dimension)) { + set_error(error, "ComfyUI quantization marker '" + marker_name + "' shape overflows"); + return false; + } + marker_elements *= static_cast(dimension); + } + if (marker_elements != marker.end - marker.begin) { + set_error(error, "ComfyUI quantization marker '" + marker_name + "' has an invalid byte length"); + return false; + } + + std::string marker_json(static_cast(marker.end - marker.begin), '\0'); + file.clear(); + file.seekg(static_cast(data_start + marker.begin)); + file.read(marker_json.data(), static_cast(marker_json.size())); + if (!file) { + set_error(error, "failed to read ComfyUI quantization marker '" + marker_name + "'"); + return false; + } + + nlohmann::json config; + try { + config = nlohmann::json::parse(marker_json); + } catch (const std::exception&) { + set_error(error, "invalid JSON in ComfyUI quantization marker '" + marker_name + "'"); + return false; + } + if (!config.is_object() || !config.contains("format") || !config["format"].is_string()) { + set_error(error, "invalid ComfyUI quantization marker '" + marker_name + "'"); + return false; + } + if (config["format"].get() != "int8_tensorwise") { + continue; + } + + const std::string base = marker_name.substr(0, marker_name.size() - marker_suffix_len); + const auto weight_it = tensors.find(base + ".weight"); + const auto scale_it = tensors.find(base + ".weight_scale"); + if (weight_it == tensors.end() || scale_it == tensors.end()) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' is missing its weight or weight_scale tensor"); + return false; + } + const SafetensorsTensorInfo& weight = weight_it->second; + const SafetensorsTensorInfo& scale = scale_it->second; + if (weight.dtype != "I8" || weight.shape.size() != 2 || scale.dtype != "F32" || + scale.shape.size() != 2 || scale.shape[0] != weight.shape[0] || scale.shape[1] != 1) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' has incompatible weight and scale tensors"); + return false; + } + const uint64_t output_rows = static_cast(weight.shape[0]); + if (output_rows > std::numeric_limits::max() / sizeof(float) || + scale.end - scale.begin != output_rows * sizeof(float)) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' has incompatible weight and scale tensors"); + return false; + } + + ComfyInt8Info info; + if (config.contains("convrot")) { + if (!config["convrot"].is_boolean()) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' has a non-boolean convrot field"); + return false; + } + info.convrot = config["convrot"].get(); + } + if (info.convrot) { + if (!config.contains("convrot_groupsize") || !config["convrot_groupsize"].is_number_unsigned()) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' has no valid ConvRot group size"); + return false; + } + uint64_t group_size = config["convrot_groupsize"].get(); + if (group_size != 256 || !is_power_of_four(group_size) || + static_cast(weight.shape[1]) % group_size != 0) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' uses an unsupported ConvRot group size"); + return false; + } + info.group_size = static_cast(group_size); + } else if (config.contains("convrot_groupsize")) { + set_error(error, "ComfyUI Int8 marker '" + marker_name + "' declares a group size without ConvRot"); + return false; + } + info.scale.name = base + ".weight_scale"; + info.scale.type = GGML_TYPE_F32; + info.scale.n_dims = 2; + // TensorStorage uses GGML's column-first shape convention. + info.scale.ne[0] = 1; + info.scale.ne[1] = weight.shape[0]; + info.scale.offset = data_start + scale.begin; + info.scale.nbytes = scale.end - scale.begin; + result->emplace(weight_it->first, info); + scale_tensor_names->emplace(scale_it->first); + } + return true; +} + // https://huggingface.co/docs/safetensors/index bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, @@ -163,29 +356,41 @@ bool read_safetensors_file(const std::string& file_path, } } - tensor_storages.clear(); - for (auto& item : header_.items()) { - std::string name = item.key(); - nlohmann::json tensor_info = item.value(); - // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str()); - - if (name == "__metadata__") { + const uint64_t data_size = file_size_ - data_start; + std::map tensor_infos; + for (const auto& item : header_.items()) { + if (item.key() == "__metadata__") { continue; } + SafetensorsTensorInfo info; + if (!read_safetensors_tensor_info(item.value(), item.key(), data_size, &info, error)) { + return false; + } + tensor_infos.emplace(item.key(), std::move(info)); + } + std::map comfy_int8_tensors; + std::set comfy_int8_scale_tensors; + if (!read_comfy_int8_metadata(file, + tensor_infos, + data_start, + &comfy_int8_tensors, + &comfy_int8_scale_tensors, + error)) { + return false; + } + + tensor_storages.clear(); + for (const auto& [name, tensor_info] : tensor_infos) { + // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str()); - std::string dtype = tensor_info["dtype"]; - nlohmann::json shape = tensor_info["shape"]; + const std::string& dtype = tensor_info.dtype; - if (dtype == "U8") { + if (dtype == "U8" || comfy_int8_scale_tensors.find(name) != comfy_int8_scale_tensors.end()) { continue; } - size_t begin = tensor_info["data_offsets"][0].get(); - size_t end = tensor_info["data_offsets"][1].get(); - if (begin > end || end > file_size_ - data_start) { - set_error(error, "data offsets out of bounds for tensor '" + name + "'"); - return false; - } + const uint64_t begin = tensor_info.begin; + const uint64_t end = tensor_info.end; ggml_type type = safetensors_dtype_to_ggml_type(dtype); if (type == GGML_TYPE_COUNT) { @@ -193,18 +398,17 @@ bool read_safetensors_file(const std::string& file_path, return false; } - if (shape.size() > SD_MAX_DIMS) { - set_error(error, "invalid tensor '" + name + "'"); - return false; - } - - int n_dims = (int)shape.size(); + int n_dims = static_cast(tensor_info.shape.size()); int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; for (int i = 0; i < n_dims; i++) { - ne[i] = shape[i].get(); + ne[i] = tensor_info.shape[i]; } if (n_dims == 5) { + if (ne[0] > std::numeric_limits::max() / ne[1]) { + set_error(error, "tensor dimensions overflow for '" + name + "'"); + return false; + } n_dims = 4; ne[0] = ne[0] * ne[1]; ne[1] = ne[2]; @@ -220,7 +424,7 @@ bool read_safetensors_file(const std::string& file_path, TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin); tensor_storage.reverse_ne(); - size_t tensor_data_size = end - begin; + uint64_t tensor_data_size = end - begin; bool tensor_size_ok; if (dtype == "F8_E4M3") { @@ -247,6 +451,20 @@ bool read_safetensors_file(const std::string& file_path, return false; } + auto comfy_int8 = comfy_int8_tensors.find(name); + if (dtype == "I8") { + if (comfy_int8 == comfy_int8_tensors.end()) { + set_error(error, "unsupported Int8 safetensors tensor '" + name + "' without a ComfyUI Int8 marker"); + return false; + } + tensor_storage.is_comfy_int8_tensorwise = true; + tensor_storage.comfy_int8_convrot = comfy_int8->second.convrot; + tensor_storage.comfy_int8_group_size = comfy_int8->second.group_size; + tensor_storage.comfy_int8_scale = comfy_int8->second.scale; + // The runtime reconstructs an F16 matrix before backend upload. + tensor_storage.expected_type = GGML_TYPE_F16; + } + tensor_storages.push_back(tensor_storage); // LOG_DEBUG("%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 5c977f516..af015a0f1 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -13,6 +13,23 @@ #define SD_MAX_DIMS 5 +// A safetensors sidecar is deliberately not added to the model parameter map. +// It remains addressable from its owning tensor so compound on-disk formats can +// upload both buffers without exposing implementation metadata as a second +// model parameter. +struct TensorStorageSidecar { + std::string name; + ggml_type type = GGML_TYPE_COUNT; + int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; + int n_dims = 0; + uint64_t offset = 0; + uint64_t nbytes = 0; + + bool valid() const { + return !name.empty() && type != GGML_TYPE_COUNT && n_dims > 0 && n_dims <= SD_MAX_DIMS && nbytes > 0; + } +}; + struct TensorStorage { std::string name; ggml_type type = GGML_TYPE_F32; @@ -22,6 +39,17 @@ struct TensorStorage { bool is_f8_e5m2 = false; bool is_f64 = false; bool is_i64 = false; + // ComfyUI TensorWiseINT8 stores the I8 weight and its per-output-row F32 + // scale separately. ConvRot metadata is carried by a U8 JSON side tensor. + // Keep the scale as an associated raw sidecar: it must not be mistaken for + // a model-level `weight_scale` parameter, but native operators need it. + bool is_comfy_int8_tensorwise = false; + bool comfy_int8_convrot = false; + // Set by the runner's backend policy before parameters are constructed. + // False selects the verified F16 compatibility reconstruction. + bool comfy_int8_native_enabled = false; + uint32_t comfy_int8_group_size = 0; + TensorStorageSidecar comfy_int8_scale; int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; int n_dims = 0; @@ -61,6 +89,15 @@ struct TensorStorage { } } + bool has_comfy_int8_scale() const { + return is_comfy_int8_tensorwise && comfy_int8_scale.valid(); + } + + bool is_comfy_int8_convrot_weight() const { + return has_comfy_int8_scale() && comfy_int8_convrot && comfy_int8_group_size == 256 && n_dims == 2 && + ne[0] > 0 && ne[1] > 0 && ne[0] % static_cast(comfy_int8_group_size) == 0; + } + void unsqueeze() { if (n_dims == 2) { n_dims = 4; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index a70ffefd3..ad83c5987 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -154,6 +155,84 @@ void i64_to_i32_vec(int64_t* src, int32_t* dst, int64_t n) { } } +static void apply_regular_hadamard_4(float* values, size_t stride) { + const float a = values[0 * stride]; + const float b = values[1 * stride]; + const float c = values[2 * stride]; + const float d = values[3 * stride]; + values[0 * stride] = (a + b + c - d) * 0.5f; + values[1 * stride] = (a + b - c + d) * 0.5f; + values[2 * stride] = (a - b + c + d) * 0.5f; + values[3 * stride] = (-a + b + c + d) * 0.5f; +} + +static bool dequantize_comfy_int8_tensorwise(const TensorStorage& tensor_storage, + const int8_t* quantized, + const float* scales, + void* dst, + ggml_type dst_type, + std::string* error) { + if (!tensor_storage.is_comfy_int8_tensorwise || tensor_storage.n_dims != 2 || + tensor_storage.ne[0] <= 0 || tensor_storage.ne[1] <= 0) { + *error = "invalid ComfyUI Int8 tensor metadata"; + return false; + } + if (dst_type != GGML_TYPE_F16 && dst_type != GGML_TYPE_F32) { + *error = "ComfyUI Int8 compatibility loading requires an F16 or F32 destination"; + return false; + } + + const size_t columns = static_cast(tensor_storage.ne[0]); + const size_t rows = static_cast(tensor_storage.ne[1]); + if (rows > std::numeric_limits::max() / columns || !tensor_storage.has_comfy_int8_scale() || + tensor_storage.comfy_int8_scale.type != GGML_TYPE_F32 || tensor_storage.comfy_int8_scale.n_dims != 2 || + tensor_storage.comfy_int8_scale.ne[0] != 1 || tensor_storage.comfy_int8_scale.ne[1] != static_cast(rows) || + tensor_storage.comfy_int8_scale.nbytes != rows * sizeof(float)) { + *error = "invalid ComfyUI Int8 tensor dimensions or scale size"; + return false; + } + const size_t group_size = tensor_storage.comfy_int8_convrot + ? static_cast(tensor_storage.comfy_int8_group_size) + : std::min(columns, 4096); + if (group_size == 0 || columns % group_size != 0) { + *error = "invalid ComfyUI Int8 ConvRot group size"; + return false; + } + + std::vector values(group_size); + for (size_t row = 0; row < rows; ++row) { + const float scale = scales[row]; + if (!std::isfinite(scale) || scale <= 0.f) { + *error = "ComfyUI Int8 tensor has a non-positive or non-finite scale"; + return false; + } + for (size_t column = 0; column < columns; column += group_size) { + const size_t offset = row * columns + column; + for (size_t i = 0; i < group_size; ++i) { + values[i] = static_cast(quantized[offset + i]) * scale; + } + if (tensor_storage.comfy_int8_convrot) { + for (size_t stride = 1; stride < group_size; stride *= 4) { + const size_t block = stride * 4; + for (size_t base = 0; base < group_size; base += block) { + for (size_t i = 0; i < stride; ++i) { + apply_regular_hadamard_4(values.data() + base + i, stride); + } + } + } + } + if (dst_type == GGML_TYPE_F16) { + auto* output = static_cast(dst) + offset; + ggml_fp32_to_fp16_row(values.data(), output, static_cast(group_size)); + } else { + auto* output = static_cast(dst) + offset; + memcpy(output, values.data(), group_size * sizeof(float)); + } + } + } + return true; +} + void convert_tensor(void* src, ggml_type src_type, void* dst, @@ -1240,10 +1319,42 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, return true; }; + auto read_comfy_int8_scales = [&](char* buf, size_t n) -> bool { + if (zip != nullptr) { + LOG_ERROR("ComfyUI Int8 tensor '%s' cannot be stored in a zip file", tensor_storage.name.c_str()); + return false; + } + if (mmapped) { + if (!mmapped->copy_data(buf, n, tensor_storage.comfy_int8_scale.offset)) { + LOG_ERROR("read ComfyUI Int8 scales failed: '%s'", file_path.c_str()); + return false; + } + } else { + file.clear(); + file.seekg(static_cast(tensor_storage.comfy_int8_scale.offset)); + file.read(buf, static_cast(n)); + if (!file) { + LOG_ERROR("read ComfyUI Int8 scales failed: '%s'", file_path.c_str()); + return false; + } + } + return true; + }; + char* read_buf = nullptr; char* target_buf = nullptr; char* convert_buf = nullptr; - if (dst_tensor->buffer == nullptr || ggml_backend_buffer_is_host(dst_tensor->buffer)) { + const bool is_comfy_int8 = tensor_storage.is_comfy_int8_tensorwise; + if (is_comfy_int8) { + read_buffer.resize(nbytes_to_read); + read_buf = reinterpret_cast(read_buffer.data()); + if (dst_tensor->buffer == nullptr || ggml_backend_buffer_is_host(dst_tensor->buffer)) { + target_buf = reinterpret_cast(dst_tensor->data); + } else { + convert_buffer.resize(ggml_nbytes(dst_tensor)); + target_buf = reinterpret_cast(convert_buffer.data()); + } + } else if (dst_tensor->buffer == nullptr || ggml_backend_buffer_is_host(dst_tensor->buffer)) { if (tensor_storage.type == dst_tensor->type) { GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes()); if (tensor_storage.is_f64 || tensor_storage.is_i64) { @@ -1279,7 +1390,27 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, read_time_ms.fetch_add(t1 - t0); t0 = ggml_time_ms(); - if (tensor_storage.is_f8_e4m3) { + if (is_comfy_int8) { + std::vector scale_buffer(tensor_storage.comfy_int8_scale.nbytes); + if (!read_comfy_int8_scales(reinterpret_cast(scale_buffer.data()), scale_buffer.size())) { + failed = true; + break; + } + std::string dequantization_error; + if (!dequantize_comfy_int8_tensorwise(tensor_storage, + reinterpret_cast(read_buf), + reinterpret_cast(scale_buffer.data()), + target_buf, + dst_tensor->type, + &dequantization_error)) { + LOG_ERROR("ComfyUI Int8 tensor '%s' cannot be reconstructed: %s", + tensor_storage.name.c_str(), + dequantization_error.c_str()); + failed = true; + break; + } + convert_buf = target_buf; + } else if (tensor_storage.is_f8_e4m3) { f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); } else if (tensor_storage.is_f8_e5m2) { f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); @@ -1288,7 +1419,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, } else if (tensor_storage.is_i64) { i64_to_i32_vec((int64_t*)read_buf, (int32_t*)target_buf, tensor_storage.nelements()); } - if (tensor_storage.type != dst_tensor->type) { + if (!is_comfy_int8 && tensor_storage.type != dst_tensor->type) { if (convert_buf == nullptr) { LOG_ERROR("read tensor data failed: too less memory for conversion"); failed = true; @@ -1303,7 +1434,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, tensor_storage.nelements() / tensor_storage.ne[0], tensor_storage.ne[0], std::move(imatrix)); - } else { + } else if (!is_comfy_int8) { convert_buf = read_buf; } t1 = ggml_time_ms(); @@ -1319,7 +1450,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, copy_to_backend_time_ms.fetch_add(t1 - t0); } - bytes_processed.fetch_add((uint64_t)nbytes_to_read); + bytes_processed.fetch_add((uint64_t)nbytes_to_read + + (is_comfy_int8 ? tensor_storage.comfy_int8_scale.nbytes : 0)); } if (zip != nullptr) { zip_close(zip); @@ -1427,6 +1559,92 @@ bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* return true; } +bool ModelLoader::load_comfy_int8_tensorwise(const TensorStorage& tensor_storage, + ggml_tensor* dst_weight, + ggml_tensor* dst_scale) { + if (!tensor_storage.is_comfy_int8_convrot_weight()) { + LOG_ERROR("native ComfyUI Int8 load requested for invalid ConvRot tensor '%s'", tensor_storage.name.c_str()); + return false; + } + if (dst_weight == nullptr || dst_scale == nullptr || dst_weight->data == nullptr || dst_scale->data == nullptr) { + LOG_ERROR("native ComfyUI Int8 load has null destination for '%s'", tensor_storage.name.c_str()); + return false; + } + if (dst_weight->type != GGML_TYPE_I8 || ggml_n_dims(dst_weight) != 2 || dst_weight->ne[0] != tensor_storage.ne[0] || + dst_weight->ne[1] != tensor_storage.ne[1] || ggml_nbytes(dst_weight) != static_cast(tensor_storage.nbytes())) { + LOG_ERROR("native ComfyUI Int8 weight destination is incompatible for '%s'", tensor_storage.name.c_str()); + return false; + } + const int64_t output_rows = tensor_storage.ne[1]; + const bool scale_shape_ok = dst_scale->type == GGML_TYPE_F32 && + ((ggml_n_dims(dst_scale) == 1 && dst_scale->ne[0] == output_rows) || + (ggml_n_dims(dst_scale) == 2 && dst_scale->ne[0] == 1 && dst_scale->ne[1] == output_rows)); + if (!scale_shape_ok || ggml_nbytes(dst_scale) != tensor_storage.comfy_int8_scale.nbytes) { + LOG_ERROR("native ComfyUI Int8 scale destination is incompatible for '%s'", tensor_storage.name.c_str()); + return false; + } + if (tensor_storage.file_index >= file_paths_.size()) { + LOG_ERROR("native ComfyUI Int8 source file is unavailable for '%s'", tensor_storage.name.c_str()); + return false; + } + if (tensor_storage.index_in_zip >= 0) { + LOG_ERROR("native ComfyUI Int8 tensor '%s' cannot be loaded from a zip container", tensor_storage.name.c_str()); + return false; + } + + const size_t weight_nbytes = static_cast(tensor_storage.nbytes()); + const size_t scale_nbytes = static_cast(tensor_storage.comfy_int8_scale.nbytes); + std::vector weights(weight_nbytes); + std::vector scales(scale_nbytes); + std::ifstream file(file_paths_[tensor_storage.file_index], std::ios::binary); + if (!file.is_open()) { + LOG_ERROR("failed to open native ComfyUI Int8 source '%s'", file_paths_[tensor_storage.file_index].c_str()); + return false; + } + const auto read_at = [&](uint64_t offset, uint8_t* dst, size_t n, const char* what) -> bool { + if (offset > static_cast(std::numeric_limits::max())) { + LOG_ERROR("native ComfyUI Int8 %s offset overflows for '%s'", what, tensor_storage.name.c_str()); + return false; + } + file.clear(); + file.seekg(static_cast(offset)); + file.read(reinterpret_cast(dst), static_cast(n)); + if (!file) { + LOG_ERROR("failed to read native ComfyUI Int8 %s for '%s'", what, tensor_storage.name.c_str()); + return false; + } + return true; + }; + if (!read_at(tensor_storage.offset, weights.data(), weights.size(), "weight") || + !read_at(tensor_storage.comfy_int8_scale.offset, scales.data(), scales.size(), "scale")) { + return false; + } + + // The compact native path does not dequantize on the host, so validate the + // sidecar before uploading it. Otherwise a malformed zero/NaN scale + // reaches the backend without the compatibility loader's validation. + const size_t scale_count = static_cast(output_rows); + for (size_t row = 0; row < scale_count; ++row) { + float scale; + memcpy(&scale, scales.data() + row * sizeof(scale), sizeof(scale)); + if (!std::isfinite(scale) || scale <= 0.f) { + LOG_ERROR("native ComfyUI Int8 tensor '%s' has a non-positive or non-finite scale", tensor_storage.name.c_str()); + return false; + } + } + + const auto upload = [](ggml_tensor* dst, const void* src, size_t n) { + if (dst->buffer != nullptr && !ggml_backend_buffer_is_host(dst->buffer)) { + ggml_backend_tensor_set(dst, src, 0, n); + } else { + memcpy(dst->data, src, n); + } + }; + upload(dst_weight, weights.data(), weights.size()); + upload(dst_scale, scales.data(), scales.size()); + return true; +} + bool ModelLoader::load_float_tensor(const std::string& name, std::vector& data, int n_threads, diff --git a/src/model_loader.h b/src/model_loader.h index f7ebcf3ef..c2e3b5bfe 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -83,6 +83,12 @@ class ModelLoader { int n_threads = 0, bool use_mmap = false); bool load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor); + // Upload the raw I8 weight and its associated F32 scale sidecar. This is + // intentionally separate from load_tensor(), whose TensorWiseINT8 path is + // the explicit F16/F32 compatibility reconstruction. + bool load_comfy_int8_tensorwise(const TensorStorage& tensor_storage, + ggml_tensor* dst_weight, + ggml_tensor* dst_scale); std::vector get_tensor_names() const { std::vector names; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 825af29b9..a2aec5cc8 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -37,6 +37,21 @@ static std::string lora_id(const ModelManager::LoraSpec& lora) { return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path; } +// `weight.convrot_scale` is an internal parameter name owned by a native +// ConvRot Linear. It maps to the scale sidecar associated with the preceding +// safetensors `weight`, never to an ordinary model `weight_scale` tensor. +static bool convrot_scale_weight_name(const std::string& name, std::string* weight_name) { + static constexpr const char* suffix = ".convrot_scale"; + static constexpr size_t suffix_len = 14; + if (!ends_with(name, suffix) || name.size() == suffix_len) { + return false; + } + if (weight_name != nullptr) { + *weight_name = name.substr(0, name.size() - suffix_len); + } + return true; +} + static bool backend_supports_host_buffer(ggml_backend_t backend) { if (backend == nullptr) { return false; @@ -680,6 +695,21 @@ bool ModelManager::validate_tensor(const TensorState& state) const { } const auto& tensor_storage_map = model_loader_.get_tensor_storage_map(); + std::string convrot_weight_name; + if (convrot_scale_weight_name(state.name, &convrot_weight_name)) { + const auto weight_it = tensor_storage_map.find(convrot_weight_name); + if (weight_it == tensor_storage_map.end() || !weight_it->second.is_comfy_int8_convrot_weight()) { + LOG_ERROR("%s ConvRot scale '%s' has no associated marked weight", state.desc.c_str(), state.name.c_str()); + return false; + } + const TensorStorage& weight = weight_it->second; + if (state.tensor->type != GGML_TYPE_F32 || ggml_n_dims(state.tensor) != 1 || + state.tensor->ne[0] != weight.ne[1] || ggml_nbytes(state.tensor) != weight.comfy_int8_scale.nbytes) { + LOG_ERROR("%s ConvRot scale '%s' has incompatible shape or type", state.desc.c_str(), state.name.c_str()); + return false; + } + return true; + } auto ts_it = tensor_storage_map.find(state.name); if (ts_it == tensor_storage_map.end()) { LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str()); @@ -747,6 +777,15 @@ bool ModelManager::can_mmap_storage(const TensorState& state) const { if (state.compute_backend == nullptr || state.params_backend == nullptr) { return false; } + std::string convrot_weight_name; + if (convrot_scale_weight_name(state.name, &convrot_weight_name)) { + return false; + } + const auto storage_it = model_loader_.get_tensor_storage_map().find(state.name); + if (storage_it != model_loader_.get_tensor_storage_map().end() && storage_it->second.is_comfy_int8_tensorwise) { + // The I8 bytes must be uploaded together with their sidecar scale. + return false; + } return sd_backend_is_cpu(state.compute_backend) || sd_backend_is_cpu(state.params_backend) || backend_supports_host_buffer(state.compute_backend); @@ -878,6 +917,31 @@ bool ModelManager::load_tensors(const std::vector& states) { return false; } + // Only the native Linear registers an I8 destination plus its private + // sidecar parameter. The explicit compatibility policy deliberately + // registers an F16 destination and must continue through load_tensor, + // where the marked bytes are reconstructed to F16. + if (tensor_storage.is_comfy_int8_convrot_weight() && state->tensor->type == GGML_TYPE_I8) { + const std::string scale_name = name + ".convrot_scale"; + const auto scale_it = states_by_name.find(scale_name); + if (scale_it == states_by_name.end() || scale_it->second == nullptr || + scale_it->second->tensor == nullptr) { + LOG_ERROR("native ConvRot tensor '%s' is missing its dedicated scale parameter", name.c_str()); + return false; + } + if (!model_loader_.load_comfy_int8_tensorwise(tensor_storage, state->tensor, scale_it->second->tensor)) { + return false; + } + { + std::lock_guard lock(loaded_names_mutex); + loaded_names.insert(name); + loaded_names.insert(scale_name); + } + // The raw uploader has already initialized both tensors. Do not + // let ModelLoader select its F16 compatibility reconstruction. + return true; + } + if (state->tensor->ne[0] != tensor_storage.ne[0] || state->tensor->ne[1] != tensor_storage.ne[1] || state->tensor->ne[2] != tensor_storage.ne[2] || diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 9db2fc217..6b3c51484 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -3995,7 +3995,15 @@ sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { return nullptr; } - if (!sd_ctx->sd->init(sd_ctx_params)) { + bool initialized = false; + try { + initialized = sd_ctx->sd->init(sd_ctx_params); + } catch (const std::exception& error) { + LOG_ERROR("failed to initialize Stable Diffusion context: %s", error.what()); + } catch (...) { + LOG_ERROR("failed to initialize Stable Diffusion context: unknown exception"); + } + if (!initialized) { delete sd_ctx->sd; sd_ctx->sd = nullptr; free(sd_ctx); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6d0572659..ec7c8ada4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -2,6 +2,14 @@ add_executable(test-ltx-rope test-ltx-rope.cpp) target_link_libraries(test-ltx-rope PRIVATE stable-diffusion) add_test(NAME test-ltx-rope COMMAND test-ltx-rope) +add_executable(test-minimax-h3-rope test_minimax_h3_rope.cpp) +target_link_libraries(test-minimax-h3-rope PRIVATE stable-diffusion) +add_test(NAME test-minimax-h3-rope COMMAND test-minimax-h3-rope) + +add_executable(test-h3-swiglu test_h3_swiglu.cpp) +target_link_libraries(test-h3-swiglu PRIVATE stable-diffusion ggml-cpu) +add_test(NAME test-h3-swiglu COMMAND test-h3-swiglu) + add_executable(test-vae-routing test-vae-routing.cpp) target_link_libraries(test-vae-routing PRIVATE stable-diffusion) add_test(NAME test-vae-routing COMMAND test-vae-routing) @@ -47,3 +55,9 @@ target_include_directories(test-gguf-comfy-shape PRIVATE "${PROJECT_SOURCE_DIR}/src") target_link_libraries(test-gguf-comfy-shape PRIVATE stable-diffusion) add_test(NAME test-gguf-comfy-shape COMMAND test-gguf-comfy-shape) + +add_executable(test-safetensors-convrot test-safetensors-convrot.cpp) +target_include_directories(test-safetensors-convrot PRIVATE + "${PROJECT_SOURCE_DIR}/src") +target_link_libraries(test-safetensors-convrot PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) +add_test(NAME test-safetensors-convrot COMMAND test-safetensors-convrot) diff --git a/tests/test-safetensors-convrot.cpp b/tests/test-safetensors-convrot.cpp new file mode 100644 index 000000000..605dba451 --- /dev/null +++ b/tests/test-safetensors-convrot.cpp @@ -0,0 +1,202 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "model_io/binary_io.h" +#include "model_loader.h" + +namespace { + +std::string make_header(const std::string& marker) { + return "{\"layer.weight\":{\"dtype\":\"I8\",\"shape\":[4,256],\"data_offsets\":[0,1024]}," + "\"layer.weight_scale\":{\"dtype\":\"F32\",\"shape\":[4,1],\"data_offsets\":[1024,1040]}," + "\"layer.comfy_quant\":{\"dtype\":\"U8\",\"shape\":[" + + std::to_string(marker.size()) + "],\"data_offsets\":[1040," + + std::to_string(1040 + marker.size()) + "]}}"; +} + +void write_fixture(const std::filesystem::path& path, const std::string& marker, float scale_value = 0.5f) { + const std::string header = make_header(marker); + std::vector weights(4 * 256, 0); + weights[0] = 2; + const float scales[] = {scale_value, scale_value, scale_value, scale_value}; + + std::ofstream file(path, std::ios::binary | std::ios::trunc); + GGML_ASSERT(file.is_open()); + model_io::write_u64(file, header.size()); + file.write(header.data(), static_cast(header.size())); + file.write(reinterpret_cast(weights.data()), static_cast(weights.size())); + file.write(reinterpret_cast(scales), sizeof(scales)); + file.write(marker.data(), static_cast(marker.size())); + GGML_ASSERT(file.good()); +} + +const TensorStorage& find_tensor(const ModelLoader& loader, const std::string& name) { + const auto& tensors = loader.get_tensor_storage_map(); + const auto it = tensors.find(name); + GGML_ASSERT(it != tensors.end()); + return it->second; +} + +int set_test_environment(const char* name, const char* value) { +#ifdef _WIN32 + return _putenv_s(name, value); +#else + return setenv(name, value, 1); +#endif +} + +int unset_test_environment(const char* name) { +#ifdef _WIN32 + return _putenv_s(name, ""); +#else + return unsetenv(name); +#endif +} + +ggml_backend_t init_cpu_backend() { + ggml_backend_load_all(); + return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); +} + +} // namespace + +int main() { + const std::filesystem::path path = std::filesystem::temp_directory_path() / + "stable-diffusion-convrot-test.safetensors"; + const std::string marker = + "{\"format\":\"int8_tensorwise\",\"convrot\":true,\"convrot_groupsize\":256}"; + write_fixture(path, marker); + + ModelLoader loader; + GGML_ASSERT(loader.init_from_file(path.string())); + const TensorStorage& weight = find_tensor(loader, "layer.weight"); + GGML_ASSERT(weight.type == GGML_TYPE_I8); + GGML_ASSERT(weight.expected_type == GGML_TYPE_F16); + GGML_ASSERT(weight.is_comfy_int8_tensorwise); + GGML_ASSERT(weight.comfy_int8_convrot); + GGML_ASSERT(weight.comfy_int8_group_size == 256); + GGML_ASSERT(weight.has_comfy_int8_scale()); + GGML_ASSERT(weight.comfy_int8_scale.name == "layer.weight_scale"); + GGML_ASSERT(weight.comfy_int8_scale.type == GGML_TYPE_F32); + GGML_ASSERT(weight.comfy_int8_scale.n_dims == 2); + GGML_ASSERT(weight.comfy_int8_scale.ne[0] == 1); + GGML_ASSERT(weight.comfy_int8_scale.ne[1] == 4); + GGML_ASSERT(weight.comfy_int8_scale.nbytes == 4 * sizeof(float)); + GGML_ASSERT(weight.is_comfy_int8_convrot_weight()); + GGML_ASSERT(loader.get_tensor_storage_map().find("layer.weight_scale") == loader.get_tensor_storage_map().end()); + + ggml_init_params params = {4096, nullptr, false}; + ggml_context* ctx = ggml_init(params); + GGML_ASSERT(ctx != nullptr); + ggml_tensor* decoded = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 4); + GGML_ASSERT(loader.load_tensor(weight, decoded)); + + std::vector values(4 * 256); + ggml_fp16_to_fp32_row(static_cast(decoded->data), values.data(), values.size()); + float energy = 0.f; + for (float value : values) { + energy += value * value; + } + // The normalized regular Hadamard matrix is orthogonal: ConvRot inversion + // preserves the squared norm of the dequantized first row. + GGML_ASSERT(std::fabs(energy - 1.f) < 0.002f); + ggml_free(ctx); + + ggml_init_params native_params = {ggml_tensor_overhead() * 2 + 1024 + 4 * sizeof(float) + 4096, nullptr, false}; + ggml_context* native_ctx = ggml_init(native_params); + GGML_ASSERT(native_ctx != nullptr); + ggml_tensor* raw_weight = ggml_new_tensor_2d(native_ctx, GGML_TYPE_I8, 256, 4); + ggml_tensor* raw_scale = ggml_new_tensor_1d(native_ctx, GGML_TYPE_F32, 4); + GGML_ASSERT(loader.load_comfy_int8_tensorwise(weight, raw_weight, raw_scale)); + GGML_ASSERT(static_cast(raw_weight->data)[0] == 2); + GGML_ASSERT(static_cast(raw_weight->data)[1] == 0); + GGML_ASSERT(static_cast(raw_scale->data)[0] == 0.5f); + GGML_ASSERT(static_cast(raw_scale->data)[3] == 0.5f); + ggml_free(native_ctx); + + ggml_backend_t cpu_backend = init_cpu_backend(); + GGML_ASSERT(cpu_backend != nullptr); + const auto native_selection = select_convrot_tensor_storage(cpu_backend, + loader.get_tensor_storage_map(), + "ConvRot loader test"); + GGML_ASSERT(native_selection.at("layer.weight").comfy_int8_native_enabled); + GGML_ASSERT(ggml_backend_supports_convrot_op(cpu_backend)); + + // A runner for another component must not inherit this component's + // ConvRot requirement when both live in the shared storage map. + const auto other_component_selection = select_convrot_tensor_storage(nullptr, + loader.get_tensor_storage_map(), + "unrelated component", + "other_component."); + GGML_ASSERT(!other_component_selection.at("layer.weight").comfy_int8_native_enabled); + + GGML_ASSERT(set_test_environment("SD_CONVROT_MODE", "compat") == 0); + const auto compatibility_selection = select_convrot_tensor_storage(cpu_backend, + loader.get_tensor_storage_map(), + "ConvRot loader test"); + GGML_ASSERT(!compatibility_selection.at("layer.weight").comfy_int8_native_enabled); + GGML_ASSERT(unset_test_environment("SD_CONVROT_MODE") == 0); + ggml_backend_free(cpu_backend); + + write_fixture(path, marker, std::numeric_limits::quiet_NaN()); + ModelLoader invalid_scale_loader; + GGML_ASSERT(invalid_scale_loader.init_from_file(path.string())); + const TensorStorage& invalid_scale_weight = find_tensor(invalid_scale_loader, "layer.weight"); + ggml_init_params invalid_scale_params = {ggml_tensor_overhead() * 2 + 1024 + 4 * sizeof(float) + 4096, nullptr, false}; + ggml_context* invalid_scale_ctx = ggml_init(invalid_scale_params); + GGML_ASSERT(invalid_scale_ctx != nullptr); + ggml_tensor* invalid_raw_weight = ggml_new_tensor_2d(invalid_scale_ctx, GGML_TYPE_I8, 256, 4); + ggml_tensor* invalid_raw_scale = ggml_new_tensor_1d(invalid_scale_ctx, GGML_TYPE_F32, 4); + GGML_ASSERT(!invalid_scale_loader.load_comfy_int8_tensorwise(invalid_scale_weight, invalid_raw_weight, invalid_raw_scale)); + ggml_free(invalid_scale_ctx); + + const std::string unsupported_group = + "{\"format\":\"int8_tensorwise\",\"convrot\":true,\"convrot_groupsize\":16}"; + write_fixture(path, unsupported_group); + ModelLoader invalid_loader; + GGML_ASSERT(!invalid_loader.init_from_file(path.string())); + + write_fixture(path, "not-json"); + ModelLoader malformed_loader; + GGML_ASSERT(!malformed_loader.init_from_file(path.string())); + + if (const char* real_model_path = std::getenv("CONVROT_MODEL_PATH")) { + ModelLoader real_loader; + GGML_ASSERT(real_loader.init_from_file(real_model_path)); + const TensorStorage* real_weight = nullptr; + for (const auto& [_, tensor] : real_loader.get_tensor_storage_map()) { + if (tensor.is_comfy_int8_tensorwise) { + real_weight = &tensor; + break; + } + } + GGML_ASSERT(real_weight != nullptr); + const size_t elements = static_cast(real_weight->ne[0]) * + static_cast(real_weight->ne[1]); + ggml_init_params real_params = {ggml_tensor_overhead() + elements * sizeof(ggml_fp16_t) + 4096, + nullptr, + false}; + ggml_context* real_ctx = ggml_init(real_params); + GGML_ASSERT(real_ctx != nullptr); + ggml_tensor* real_decoded = ggml_new_tensor_2d(real_ctx, + GGML_TYPE_F16, + real_weight->ne[0], + real_weight->ne[1]); + GGML_ASSERT(real_loader.load_tensor(*real_weight, real_decoded)); + const float first_value = ggml_fp16_to_fp32(static_cast(real_decoded->data)[0]); + GGML_ASSERT(std::isfinite(first_value)); + ggml_free(real_ctx); + } + + std::error_code ec; + std::filesystem::remove(path, ec); + return 0; +} diff --git a/tests/test_h3_swiglu.cpp b/tests/test_h3_swiglu.cpp new file mode 100644 index 000000000..e2bd6810c --- /dev/null +++ b/tests/test_h3_swiglu.cpp @@ -0,0 +1,65 @@ +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include + +static bool run_case(int64_t features, int64_t columns) { + ggml_context* ctx = ggml_init({ 4 * 1024 * 1024, nullptr, true }); + if (!ctx) { + return false; + } + + ggml_tensor* input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, features * 2, columns); + ggml_tensor* first = ggml_cont(ctx, ggml_view_2d(ctx, input, features, columns, input->nb[1], 0)); + ggml_tensor* second = ggml_cont(ctx, ggml_view_2d(ctx, input, features, columns, input->nb[1], features * input->nb[0])); + ggml_tensor* old_path = ggml_mul(ctx, ggml_silu(ctx, first), second); + ggml_tensor* fused = ggml_swiglu(ctx, input); + + ggml_cgraph* graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, old_path); + ggml_build_forward_expand(graph, fused); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + ggml_backend_free(backend); + ggml_free(ctx); + return false; + } + + std::vector values(static_cast(features * 2 * columns)); + for (size_t i = 0; i < values.size(); ++i) { + values[i] = static_cast((static_cast(i * 17 % 257) - 128)) / 19.0f; + } + ggml_backend_tensor_set(input, values.data(), 0, ggml_nbytes(input)); + bool passed = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + std::vector expected(static_cast(features * columns)); + std::vector actual(static_cast(features * columns)); + if (passed) { + ggml_backend_tensor_get(old_path, expected.data(), 0, ggml_nbytes(old_path)); + ggml_backend_tensor_get(fused, actual.data(), 0, ggml_nbytes(fused)); + for (size_t i = 0; i < actual.size(); ++i) { + if (!std::isfinite(actual[i]) || std::fabs(actual[i] - expected[i]) > 2e-6f * (1.0f + std::fabs(expected[i]))) { + std::fprintf(stderr, "SwiGLU mismatch at %zu: %.9g vs %.9g\n", i, actual[i], expected[i]); + passed = false; + break; + } + } + } + + ggml_backend_buffer_free(buffer); + ggml_backend_free(backend); + ggml_free(ctx); + return passed; +} + +int main() { + const bool passed = run_case(37, 1) && run_case(14336, 23) && run_case(14336, 257); + if (passed) { + std::puts("H3 fused SwiGLU test passed"); + } + return passed ? 0 : 1; +} diff --git a/tests/test_minimax_h3_rope.cpp b/tests/test_minimax_h3_rope.cpp new file mode 100644 index 000000000..ed919e322 --- /dev/null +++ b/tests/test_minimax_h3_rope.cpp @@ -0,0 +1,291 @@ +#include "model/diffusion/minimax_h3.hpp" + +#include +#include +#include +#include +#include +#include +#include + +enum class RopePattern { IDENTITY, QUARTER_TURN, GENERAL }; + +struct RopeCase { + int64_t dimensions; + int64_t rotated; + int64_t heads; + int64_t tokens; + int64_t batches; + bool strided; + RopePattern pattern; +}; + +static size_t input_index(const RopeCase& shape, int64_t d, int64_t h, int64_t l, int64_t n) { + const int64_t width = shape.dimensions + (shape.strided ? 8 : 0); + return static_cast(d + (shape.strided ? 3 : 0) + width * (h + shape.heads * (l + shape.tokens * n))); +} + +static bool supports_fused_rope(ggml_backend_t backend, const RopeCase& shape) { + ggml_init_params params = {1024 * 1024, nullptr, true}; + ggml_context* ctx = ggml_init(params); + if (!ctx) { + return false; + } + auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, shape.rotated, shape.heads, shape.tokens, shape.batches); + auto pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, shape.rotated / 2, shape.tokens); + const bool supported = ggml_backend_supports_op(backend, ggml_rope_flux(ctx, x, pe)); + ggml_free(ctx); + return supported; +} + +static bool run_path(ggml_backend_t backend, + const RopeCase& shape, + const std::vector& input, + const std::vector& positions, + bool use_fused, + bool expect_fused, + bool benchmark, + std::vector* actual) { + ggml_init_params params = {8 * 1024 * 1024, nullptr, true}; + ggml_context* ctx = ggml_init(params); + if (!ctx) { + return false; + } + const int64_t width = shape.dimensions + (shape.strided ? 8 : 0); + auto storage = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, shape.heads, shape.tokens, shape.batches); + auto x = shape.strided ? ggml_view_4d(ctx, storage, shape.dimensions, shape.heads, shape.tokens, shape.batches, + storage->nb[1], storage->nb[2], storage->nb[3], 3 * sizeof(float)) + : storage; + auto pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, shape.rotated / 2, shape.tokens); + auto result = MiniMaxH3::apply_partial_rope(ctx, x, pe, use_fused ? backend : nullptr); + auto graph = ggml_new_graph_custom(ctx, 2048, false); + ggml_build_forward_expand(graph, result); + int fused_nodes = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + fused_nodes += ggml_graph_node(graph, i)->op == GGML_OP_ROPE_FLUX; + } + if ((fused_nodes != 0) != expect_fused) { + std::fprintf(stderr, "unexpected RoPE graph: fused nodes=%d expected=%d\n", fused_nodes, expect_fused); + ggml_free(ctx); + return false; + } + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + std::fprintf(stderr, "could not allocate RoPE test tensors\n"); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(storage, input.data(), 0, ggml_nbytes(storage)); + ggml_backend_tensor_set(pe, positions.data(), 0, ggml_nbytes(pe)); + bool passed = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (passed && benchmark) { + constexpr int runs = 3; + const auto start = std::chrono::steady_clock::now(); + for (int i = 0; i < runs; ++i) { + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + passed = false; + break; + } + } + ggml_backend_synchronize(backend); + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - start).count() / runs; + std::printf("H3 partial RoPE backend=%s path=%s D=%lld R=%lld H=%lld L=%lld N=%lld %.3f ms/run\n", + ggml_backend_name(backend), expect_fused ? "fused" : "fallback", + static_cast(shape.dimensions), static_cast(shape.rotated), + static_cast(shape.heads), static_cast(shape.tokens), + static_cast(shape.batches), ms); + } + if (passed) { + actual->resize(ggml_nelements(result)); + ggml_backend_tensor_get(result, actual->data(), 0, ggml_nbytes(result)); + } + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return passed; +} + +static bool run_case(ggml_backend_t backend, const RopeCase& shape, bool benchmark = false) { + const int64_t width = shape.dimensions + (shape.strided ? 8 : 0); + std::vector input(width * shape.heads * shape.tokens * shape.batches, 5000.0f); + std::vector positions(2 * shape.rotated * shape.tokens); + for (int64_t n = 0; n < shape.batches; ++n) { + for (int64_t l = 0; l < shape.tokens; ++l) { + for (int64_t h = 0; h < shape.heads; ++h) { + for (int64_t d = 0; d < shape.dimensions; ++d) { + input[input_index(shape, d, h, l, n)] = d == shape.rotated ? -0.0f : + static_cast((d * 11 + h * 7 + l * 13 + n * 17) % 251 - 125) * 0.015625f; + } + } + } + } + for (int64_t l = 0; l < shape.tokens; ++l) { + for (int64_t p = 0; p < shape.rotated / 2; ++p) { + float* matrix = positions.data() + 4 * (p + (shape.rotated / 2) * l); + if (shape.pattern == RopePattern::IDENTITY) { + matrix[0] = matrix[3] = 1.0f; + matrix[1] = matrix[2] = 0.0f; + } else if (shape.pattern == RopePattern::QUARTER_TURN) { + matrix[0] = matrix[3] = 0.0f; + matrix[1] = -1.0f; + matrix[2] = 1.0f; + } else { + // Non-symmetric matrices expose row/column swaps that identity + // rotations, or accidentally paired inverse rotations, hide. + const double angle = (p * 7 + l * 3) * 0.031; + matrix[0] = static_cast(0.81 + 0.17 * std::sin(angle)); + matrix[1] = static_cast(-0.31 + 0.11 * std::cos(angle)); + matrix[2] = static_cast(0.57 + 0.13 * std::sin(angle + 0.2)); + matrix[3] = static_cast(0.73 - 0.19 * std::cos(angle + 0.3)); + } + } + } + + const bool fused_supported = supports_fused_rope(backend, shape) && std::getenv("GGML_ROPE_FLUX_DISABLE") == nullptr; + std::vector fallback; + std::vector selected; + if (!run_path(backend, shape, input, positions, false, false, benchmark, &fallback) || + !run_path(backend, shape, input, positions, true, fused_supported, benchmark, &selected)) { + return false; + } + double max_error = 0; + for (int64_t n = 0; n < shape.batches; ++n) { + for (int64_t h = 0; h < shape.heads; ++h) { + for (int64_t l = 0; l < shape.tokens; ++l) { + for (int64_t d = 0; d < shape.dimensions; ++d) { + const size_t index = d + shape.dimensions * (l + shape.tokens * (h + shape.heads * n)); + const float original = input[input_index(shape, d, h, l, n)]; + double expected = original; + if (d < shape.rotated) { + const int64_t pair = d % (shape.rotated / 2); + const int64_t component = d / (shape.rotated / 2); + const float* matrix = positions.data() + 4 * (pair + (shape.rotated / 2) * l); + expected = static_cast(input[input_index(shape, pair, h, l, n)]) * matrix[component * 2] + + static_cast(input[input_index(shape, pair + shape.rotated / 2, h, l, n)]) * matrix[component * 2 + 1]; + } + for (const auto* output : {&fallback, &selected}) { + const float value = (*output)[index]; + const double error = std::abs(static_cast(value) - expected); + max_error = std::max(max_error, error); + if (!std::isfinite(value) || error > 1e-5 * (1.0 + std::abs(expected)) || + (d >= shape.rotated && std::memcmp(&value, &original, sizeof(float)) != 0)) { + std::fprintf(stderr, "H3 RoPE mismatch path=%s R=%lld strided=%d index=%zu: got %.9g expected %.9g\n", + output == &fallback ? "fallback" : "selected", static_cast(shape.rotated), + shape.strided, index, value, expected); + return false; + } + } + } + } + } + } + std::printf("H3 RoPE passed backend=%s R=%lld pattern=%d strided=%d fused=%d max_abs=%.3g\n", + ggml_backend_name(backend), static_cast(shape.rotated), static_cast(shape.pattern), + shape.strided, fused_supported, max_error); + return true; +} + +// Direct operator coverage complements H3's split-half adapter: arbitrary +// element/row strides, offset views, head/batch mapping, and null-PE permutation. +static bool run_direct_flux_case(ggml_backend_t backend, bool strided, bool null_pe) { + constexpr int D = 14, H = 3, L = 7, B = 2; + ggml_context* ctx = ggml_init({1024 * 1024, nullptr, true}); + if (!ctx) return false; + auto storage = strided ? ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, D + 1, H + 1, (L + 1) * B) + : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, L, B); + auto x = storage; + size_t x_offset = 0; + if (strided) { + x_offset = sizeof(float); + x = ggml_view_4d(ctx, storage, D, H, L, B, storage->nb[2], storage->nb[3], + storage->nb[3] * (L + 1), x_offset); + x->nb[0] = 2 * sizeof(float); + } + auto pe_storage = strided ? ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 3, D / 2 + 1, L + 1) + : ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, D / 2, L); + auto pe = pe_storage; + size_t pe_offset = 0; + if (strided) { + pe_offset = sizeof(float); + pe = ggml_view_4d(ctx, pe_storage, 2, 2, D / 2, L, pe_storage->nb[1], pe_storage->nb[2], + pe_storage->nb[3], pe_offset); + pe->nb[0] = 2 * sizeof(float); + } + auto result = ggml_rope_flux(ctx, x, null_pe ? nullptr : pe); + if (!ggml_backend_supports_op(backend, result)) { + std::printf("Direct ROPE_FLUX unsupported backend=%s strided=%d null_pe=%d (skipped)\n", + ggml_backend_name(backend), strided, null_pe); + ggml_free(ctx); + return true; + } + auto graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, result); + auto buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { ggml_free(ctx); return false; } + std::vector input(ggml_nelements(storage), 5000.0f); + std::vector positions(ggml_nelements(pe_storage), 7000.0f); + auto xi = [&](int d, int h, int l, int b) { + return (x_offset + d * x->nb[0] + h * x->nb[1] + l * x->nb[2] + b * x->nb[3]) / sizeof(float); + }; + auto pi = [&](int c, int out, int p, int l) { + return (pe_offset + c * pe->nb[0] + out * pe->nb[1] + p * pe->nb[2] + l * pe->nb[3]) / sizeof(float); + }; + for (int b = 0; b < B; ++b) for (int l = 0; l < L; ++l) for (int h = 0; h < H; ++h) + for (int d = 0; d < D; ++d) input[xi(d,h,l,b)] = ((d*11+h*7+l*13+b*17)%67-33)*0.03125f; + for (int l = 0; l < L; ++l) for (int p = 0; p < D / 2; ++p) + for (int out = 0; out < 2; ++out) for (int c = 0; c < 2; ++c) + positions[pi(c,out,p,l)] = (c*3-out*5+p*7+l*11-23)*0.015625f; + ggml_backend_tensor_set(storage, input.data(), 0, ggml_nbytes(storage)); + ggml_backend_tensor_set(pe_storage, positions.data(), 0, ggml_nbytes(pe_storage)); + bool passed = true; + std::vector actual(ggml_nelements(result)); + for (int repeat = 0; repeat < 2 && passed; ++repeat) { + passed = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (!passed) break; + ggml_backend_tensor_get(result, actual.data(), 0, ggml_nbytes(result)); + for (int b = 0; b < B; ++b) for (int h = 0; h < H; ++h) for (int l = 0; l < L; ++l) + for (int d = 0; d < D; ++d) { + const int p = d / 2, out = d % 2; + const float expected = null_pe ? input[xi(d,h,l,b)] : + input[xi(2*p,h,l,b)] * positions[pi(0,out,p,l)] + + input[xi(2*p+1,h,l,b)] * positions[pi(1,out,p,l)]; + const float got = actual[d + D * (l + L * (h + H*b))]; + if (!std::isfinite(got) || std::fabs(got-expected) > 1e-6f * (1+std::fabs(expected))) { + std::fprintf(stderr, "Direct ROPE_FLUX mismatch strided=%d null_pe=%d d=%d h=%d l=%d b=%d got=%g expected=%g\n", + strided, null_pe, d,h,l,b,got,expected); + passed = false; + } + } + } + std::printf("Direct ROPE_FLUX backend=%s strided=%d null_pe=%d %s\n", ggml_backend_name(backend), + strided, null_pe, passed ? "passed" : "FAILED"); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return passed; +} + +int main(int argc, char** argv) { + ggml_backend_load_all(); + ggml_backend_t backend = ggml_backend_init_by_name(argc > 1 ? argv[1] : "CPU", nullptr); + if (!backend) { + std::fprintf(stderr, "requested RoPE backend is not available\n"); + return 1; + } + bool passed = true; + for (bool strided : {false, true}) for (bool null_pe : {false, true}) + passed = run_direct_flux_case(backend, strided, null_pe) && passed; + for (int64_t rotated : {96, 128, 32}) { + for (auto pattern : {RopePattern::IDENTITY, RopePattern::QUARTER_TURN, RopePattern::GENERAL}) { + for (bool strided : {false, true}) { + passed = run_case(backend, {128, rotated, 3, 7, 2, strided, pattern}) && passed; + } + } + } + if (std::getenv("SD_MINIMAX_H3_ROPE_BENCH")) { + passed = run_case(backend, {128, 96, 56, 15422, 1, true, RopePattern::GENERAL}, true) && passed; + } + ggml_backend_free(backend); + if (passed) { + std::puts("MiniMax H3 partial RoPE tests passed"); + } + return passed ? 0 : 1; +}