fix: apply fp8 weight scale factors when loading scaled safetensors

scripts/convert_fp8_scale_to_bf16.py

@@ -58,11 +58,28 @@ def numel(shape):
     return math.prod(shape) if shape else 1
 
 
-def scale_key_for_weight(name: str):
+def scale_keys_for_weight(name: str):
+    # Both spellings ship in the wild.
+    keys = []
     if name.endswith(".weight"):
-        return name[:-len(".weight")] + ".weight_scale"
+        base = name[:-len(".weight")]
+        keys.append(base + ".weight_scale")
+        keys.append(base + ".scale_weight")
     if name.endswith("weight"):
-        return name + "_scale"
+        keys.append(name + "_scale")
+    return keys
+
+
+def resolve_scale_key(name: str, entries):
+    for key in scale_keys_for_weight(name):
+        if key in entries:
+            return key
+    return None
+
+
+def input_scale_key_for_weight(name: str):
+    if name.endswith(".weight"):
+        return name[:-len(".weight")] + ".scale_input"
     return None
 
 
@@ -76,19 +93,23 @@ def build_output_plan(header):
     plan = []
 
     for name, info in entries.items():
-        scale_key = scale_key_for_weight(name)
-        if info["dtype"] in FP8_DTYPES and scale_key in entries:
+        scale_key = resolve_scale_key(name, entries)
+        if info["dtype"] in FP8_DTYPES and scale_key is not None:
             paired_scale_keys.add(scale_key)
+            # ".scale_input" is meaningless once the weight is materialised as BF16.
+            input_key = input_scale_key_for_weight(name)
+            if input_key is not None and input_key in entries:
+                paired_scale_keys.add(input_key)
 
     for name, info in entries.items():
         if name in paired_scale_keys:
             continue
 
         dtype = info["dtype"]
         shape = info["shape"]
-        scale_key = scale_key_for_weight(name)
+        scale_key = resolve_scale_key(name, entries)
 
-        if dtype in FP8_DTYPES and scale_key in entries:
+        if dtype in FP8_DTYPES and scale_key is not None:
             scale_info = entries[scale_key]
             plan.append(
                 {

src/model_io/safetensors_io.cpp

@@ -1,6 +1,7 @@
 #include "safetensors_io.h"
 
 #include <algorithm>
+#include <cmath>
 #include <cstdint>
 #include <exception>
 #include <filesystem>
@@ -233,6 +234,52 @@ bool read_safetensors_file(const std::string& file_path,
         comfy_quant_configs.emplace(module_name, std::move(config));
     }
 
+    std::unordered_map<std::string, float> fp8_scale_weights;
+    std::unordered_set<std::string> fp8_scale_tensor_names;
+    for (const auto& item : header_.items()) {
+        const std::string& name = item.key();
+        // Both spellings ship in the wild. Only F8 weights are paired, so int8_tensorwise
+        // checkpoints keep the existing ".weight_scale" handling below.
+        std::string suffix;
+        if (ends_with(name, ".scale_weight")) {
+            suffix = ".scale_weight";
+        } else if (ends_with(name, ".weight_scale")) {
+            suffix = ".weight_scale";
+        } else {
+            continue;
+        }
+        if (name == "__metadata__") {
+            continue;
+        }
+        const std::string module_name = name.substr(0, name.size() - suffix.size());
+        auto weight_it                = header_.find(module_name + ".weight");
+        if (weight_it == header_.end() || weight_it.value().value("dtype", "") != "F8_E4M3") {
+            continue;
+        }
+        const nlohmann::json& scale_info = item.value();
+        if (scale_info.value("dtype", "") != "F32") {
+            continue;
+        }
+        const size_t sbegin = scale_info["data_offsets"][0].get<size_t>();
+        const size_t send   = scale_info["data_offsets"][1].get<size_t>();
+        if (sbegin > send || send - sbegin != sizeof(float) || send > file_size_ - data_start) {
+            continue;
+        }
+        float scale = 1.0f;
+        file.clear();
+        file.seekg((std::streamoff)(data_start + sbegin), std::ios::beg);
+        file.read((char*)&scale, sizeof(float));
+        if (!file || !std::isfinite(scale) || scale == 0.0f) {
+            continue;
+        }
+        fp8_scale_weights[module_name] = scale;
+        fp8_scale_tensor_names.insert(name);
+        fp8_scale_tensor_names.insert(module_name + ".scale_input");
+    }
+    if (!fp8_scale_weights.empty()) {
+        LOG_DEBUG("safetensors: applying %zu fp8 scale_weight factors", fp8_scale_weights.size());
+    }
+
     tensor_storages.clear();
     for (auto& item : header_.items()) {
         std::string name           = item.key();
@@ -250,6 +297,10 @@ bool read_safetensors_file(const std::string& file_path,
             continue;
         }
 
+        if (fp8_scale_tensor_names.count(name) > 0) {
+            continue;
+        }
+
         size_t begin = tensor_info["data_offsets"][0].get<size_t>();
         size_t end   = tensor_info["data_offsets"][1].get<size_t>();
         if (begin > end || end > file_size_ - data_start) {
@@ -328,6 +379,12 @@ bool read_safetensors_file(const std::string& file_path,
         bool tensor_size_ok;
         if (dtype == "F8_E4M3") {
             tensor_storage.is_f8_e4m3 = true;
+            if (ends_with(name, ".weight")) {
+                auto scale_it = fp8_scale_weights.find(name.substr(0, name.size() - std::string(".weight").size()));
+                if (scale_it != fp8_scale_weights.end()) {
+                    tensor_storage.fp8_scale = scale_it->second;
+                }
+            }
             // f8 -> f16
             tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
         } else if (dtype == "F8_E5M2") {

src/model_io/tensor_storage.h

@@ -18,6 +18,7 @@ struct TensorStorage {
     ggml_type type              = GGML_TYPE_F32;
     ggml_type expected_type     = GGML_TYPE_COUNT;
     bool is_f8_e4m3             = false;
+    float fp8_scale             = 1.0f;  // companion scale tensor; 1.0f when absent
     bool is_f8_e5m2             = false;
     bool is_f64                 = false;
     bool is_i64                 = false;

src/model_loader.cpp

@@ -131,6 +131,12 @@ void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
     }
 }
 
+void f16_scale_vec(uint16_t* data, int64_t n, float scale) {
+    for (int64_t i = 0; i < n; i++) {
+        data[i] = ggml_fp32_to_fp16(ggml_fp16_to_fp32(data[i]) * scale);
+    }
+}
+
 void f8_e5m2_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
     // support inplace op
     for (int64_t i = n - 1; i >= 0; i--) {
@@ -1217,6 +1223,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     t0 = ggml_time_ms();
                     if (tensor_storage.is_f8_e4m3) {
                         f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
+                        if (tensor_storage.fp8_scale != 1.0f) {
+                            f16_scale_vec((uint16_t*)target_buf, tensor_storage.nelements(), tensor_storage.fp8_scale);
+                        }
                     } else if (tensor_storage.is_f8_e5m2) {
                         f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
                     } else if (tensor_storage.is_f64) {