Skip to content

feat(minimax-h3): load ComfyUI NVFP4 AWQ text encoder #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aegioscy wants to merge 1 commit into
base: 2026-08-11
Choose a base branch
from
Draft

feat(minimax-h3): load ComfyUI NVFP4 AWQ text encoder #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 29 additions & 1 deletion src/core/ggml_extend.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4130,6 +4130,8 @@ class Linear : public UnaryBlock {
bool force_prec_f32;
bool allow_weight_scale;
bool has_weight_scale = false;
bool has_nvfp4_global_scale = false;
bool has_nvfp4_pre_quant_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.
Expand All @@ -4144,6 +4146,8 @@ class Linear : public UnaryBlock {
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
this->prefix = prefix;
has_weight_scale = false;
has_nvfp4_global_scale = false;
has_nvfp4_pre_quant_scale = false;
has_convrot_weight = false;
use_convrot_f16_compat = false;
use_convrot_q8_decomp = false;
Expand Down Expand Up @@ -4180,6 +4184,21 @@ class Linear : public UnaryBlock {
wtype = GGML_TYPE_F32;
}
params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
if (storage_it != tensor_storage_map.end() && storage_it->second.is_comfy_nvfp4_weight()) {
params["nvfp4_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
has_nvfp4_global_scale = true;
const auto pre_quant_it = tensor_storage_map.find(prefix + "pre_quant_scale");
if (pre_quant_it != tensor_storage_map.end()) {
if (pre_quant_it->second.n_dims != 1 || pre_quant_it->second.ne[0] != in_features ||
(pre_quant_it->second.type != GGML_TYPE_F32 &&
pre_quant_it->second.type != GGML_TYPE_F16 &&
pre_quant_it->second.type != GGML_TYPE_BF16)) {
throw std::runtime_error("invalid ComfyUI NVFP4 pre_quant_scale shape");
}
params["pre_quant_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_features);
has_nvfp4_pre_quant_scale = true;
}
}
if (bias) {
enum ggml_type wtype = GGML_TYPE_F32;
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features);
Expand Down Expand Up @@ -4215,12 +4234,15 @@ class Linear : public UnaryBlock {
}

ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
if (has_nvfp4_pre_quant_scale) {
x = ggml_mul(ctx->ggml_ctx, x, params["pre_quant_scale"]);
}
ggml_tensor* w = params["weight"];
ggml_tensor* b = nullptr;
if (bias) {
b = params["bias"];
}
ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
ggml_tensor* linear_bias = (has_weight_scale || has_nvfp4_global_scale) ? nullptr : b;
ggml_tensor* out = nullptr;
if (has_convrot_weight) {
if (use_convrot_q8_decomp) {
Expand Down Expand Up @@ -4291,6 +4313,12 @@ class Linear : public UnaryBlock {
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
}
if (has_nvfp4_global_scale) {
out = ggml_mul(ctx->ggml_ctx, out, params["nvfp4_scale"]);
if (b != nullptr) {
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
}
return out;
}
};
Expand Down
123 changes: 120 additions & 3 deletions src/model_io/safetensors_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,11 @@ struct ComfyInt8Info {
TensorStorageSidecar scale;
};

struct ComfyNvfp4Info {
TensorStorageSidecar block_scale;
uint64_t global_scale_offset = 0;
};

static bool read_safetensors_tensor_info(const nlohmann::json& value,
const std::string& name,
uint64_t data_start,
Expand Down Expand Up @@ -307,6 +312,94 @@ static bool read_comfy_int8_metadata(std::ifstream& file,
return true;
}

static bool read_comfy_nvfp4_metadata(std::ifstream& file,
const std::map<std::string, SafetensorsTensorInfo>& tensors,
uint64_t data_start,
std::map<std::string, ComfyNvfp4Info>* result,
std::set<std::string>* sidecar_names,
std::string* error) {
result->clear();
sidecar_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" || marker.end - marker.begin > 4096) {
set_error(error, "invalid ComfyUI quantization marker '" + marker_name + "'");
return false;
}
std::string marker_json(static_cast<size_t>(marker.end - marker.begin), '\0');
file.clear();
file.seekg(static_cast<std::streamoff>(data_start + marker.begin));
file.read(marker_json.data(), static_cast<std::streamsize>(marker_json.size()));
if (!file) {
set_error(error, "failed to read ComfyUI quantization marker '" + marker_name + "'");
return false;
}
const nlohmann::json config = nlohmann::json::parse(marker_json, nullptr, false);
if (config.is_discarded() || !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<std::string>() != "nvfp4") {
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 block_it = tensors.find(base + ".weight_scale");
const auto global_it = tensors.find(base + ".weight_scale_2");
if (weight_it == tensors.end() || block_it == tensors.end() || global_it == tensors.end()) {
set_error(error, "ComfyUI NVFP4 marker '" + marker_name + "' is missing its weight or scales");
return false;
}
const auto& weight = weight_it->second;
const auto& block = block_it->second;
const auto& global = global_it->second;
if (weight.dtype != "U8" || weight.shape.size() != 2 || weight.shape[0] <= 0 ||
weight.shape[1] <= 0 || weight.shape[1] > std::numeric_limits<int64_t>::max() / 2 ||
block.dtype != "F8_E4M3" || block.shape.size() != 2 ||
global.dtype != "F32" || !global.shape.empty() || global.end - global.begin != sizeof(float)) {
set_error(error, "ComfyUI NVFP4 marker '" + marker_name + "' has incompatible tensor types");
return false;
}
const uint64_t rows = static_cast<uint64_t>(weight.shape[0]);
const uint64_t columns = 2 * static_cast<uint64_t>(weight.shape[1]);
if (columns % 64 != 0 || rows > std::numeric_limits<uint64_t>::max() - 127 ||
rows > std::numeric_limits<uint64_t>::max() / columns) {
set_error(error, "ComfyUI NVFP4 marker '" + marker_name + "' has unsupported dimensions");
return false;
}
const uint64_t padded_rows = ((rows + 127) / 128) * 128;
const uint64_t blocks_per_row = columns / 16;
const uint64_t padded_blocks = ((blocks_per_row + 3) / 4) * 4;
if (padded_rows > std::numeric_limits<uint64_t>::max() / padded_blocks ||
block.shape[0] != static_cast<int64_t>(padded_rows) ||
block.shape[1] != static_cast<int64_t>(padded_blocks) ||
block.end - block.begin != padded_rows * padded_blocks ||
weight.end - weight.begin != rows * columns / 2) {
set_error(error, "ComfyUI NVFP4 marker '" + marker_name + "' has incompatible tensor dimensions");
return false;
}
ComfyNvfp4Info info;
info.block_scale.name = block_it->first;
info.block_scale.type = GGML_TYPE_I8; // Raw E4M3 bytes, not a model parameter.
info.block_scale.n_dims = 2;
info.block_scale.ne[0] = static_cast<int64_t>(padded_blocks);
info.block_scale.ne[1] = static_cast<int64_t>(padded_rows);
info.block_scale.offset = data_start + block.begin;
info.block_scale.nbytes = block.end - block.begin;
info.global_scale_offset = data_start + global.begin;
result->emplace(weight_it->first, info);
sidecar_names->emplace(block_it->first);
sidecar_names->emplace(global_it->first);
}
return true;
}

// https://huggingface.co/docs/safetensors/index
bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
Expand Down Expand Up @@ -397,21 +490,31 @@ bool read_safetensors_file(const std::string& file_path,
error)) {
return false;
}
std::map<std::string, ComfyNvfp4Info> comfy_nvfp4_tensors;
std::set<std::string> comfy_nvfp4_sidecars;
if (!read_comfy_nvfp4_metadata(file, tensor_infos, data_start,
&comfy_nvfp4_tensors, &comfy_nvfp4_sidecars, 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());

const std::string& dtype = tensor_info.dtype;

if (dtype == "U8" || comfy_int8_scale_tensors.find(name) != comfy_int8_scale_tensors.end()) {
const auto comfy_nvfp4 = comfy_nvfp4_tensors.find(name);
if ((dtype == "U8" && comfy_nvfp4 == comfy_nvfp4_tensors.end()) ||
comfy_int8_scale_tensors.find(name) != comfy_int8_scale_tensors.end() ||
comfy_nvfp4_sidecars.find(name) != comfy_nvfp4_sidecars.end()) {
continue;
}

const uint64_t begin = tensor_info.begin;
const uint64_t end = tensor_info.end;

ggml_type type = safetensors_dtype_to_ggml_type(dtype);
ggml_type type = comfy_nvfp4 != comfy_nvfp4_tensors.end()
? GGML_TYPE_NVFP4 : safetensors_dtype_to_ggml_type(dtype);
if (type == GGML_TYPE_COUNT) {
set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
return false;
Expand All @@ -433,6 +536,9 @@ bool read_safetensors_file(const std::string& file_path,
nelements *= std::max<uint64_t>(dim, 1);
ne[i] = static_cast<int64_t>(dim);
}
if (comfy_nvfp4 != comfy_nvfp4_tensors.end()) {
ne[1] *= 2; // The last safetensors dimension packs two FP4 values per byte.
}

if (n_dims == 5) {
if (ne[1] == 0 || ne[0] > std::numeric_limits<int64_t>::max() / ne[1]) {
Expand All @@ -457,7 +563,11 @@ bool read_safetensors_file(const std::string& file_path,
uint64_t tensor_data_size = end - begin;

bool tensor_size_ok;
if (dtype == "F8_E4M3") {
if (comfy_nvfp4 != comfy_nvfp4_tensors.end()) {
tensor_storage.is_comfy_nvfp4 = true;
tensor_storage.comfy_nvfp4_block_scale = comfy_nvfp4->second.block_scale;
tensor_size_ok = (tensor_storage.nbytes_to_read() == tensor_data_size);
} else if (dtype == "F8_E4M3") {
tensor_storage.is_f8_e4m3 = true;
// f8 -> f16
tensor_size_ok = (tensor_storage.nbytes() / 2 == tensor_data_size);
Expand Down Expand Up @@ -500,6 +610,13 @@ bool read_safetensors_file(const std::string& file_path,
// LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str());
}

for (const auto& [weight_name, info] : comfy_nvfp4_tensors) {
int64_t scalar_shape[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
const std::string layer_name = weight_name.substr(0, weight_name.size() - 7);
tensor_storages.emplace_back(layer_name + ".nvfp4_scale", GGML_TYPE_F32,
scalar_shape, 1, 0, info.global_scale_offset);
}

return true;
}

Expand Down
14 changes: 13 additions & 1 deletion src/model_io/tensor_storage.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ struct TensorStorage {
bool comfy_int8_convrot_op_enabled = false;
uint32_t comfy_int8_group_size = 0;
TensorStorageSidecar comfy_int8_scale;
// ComfyUI NVFP4 stores two E2M1 values per byte. Its FP8 block scales
// are kept separately in the cuBLAS blocked layout, and the global F32
// scale is exposed as a small model parameter for the Linear graph.
bool is_comfy_nvfp4 = false;
TensorStorageSidecar comfy_nvfp4_block_scale;
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
int n_dims = 0;

Expand Down Expand Up @@ -86,7 +91,9 @@ struct TensorStorage {
}

int64_t nbytes_to_read() const {
if (is_f8_e4m3 || is_f8_e5m2) {
if (is_comfy_nvfp4) {
return nelements() / 2;
} else if (is_f8_e4m3 || is_f8_e5m2) {
return nbytes() / 2;
} else if (is_f64 || is_i64) {
return nbytes() * 2;
Expand All @@ -99,6 +106,11 @@ struct TensorStorage {
return is_comfy_int8_tensorwise && comfy_int8_scale.valid();
}

bool is_comfy_nvfp4_weight() const {
return is_comfy_nvfp4 && comfy_nvfp4_block_scale.valid() && n_dims == 2 &&
ne[0] > 0 && ne[1] > 0 && ne[0] % 64 == 0;
}

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<int64_t>(comfy_int8_group_size) == 0;
Expand Down
Loading
Loading