[0.17] Load FP8 weights as F16 on Metal instead of aborting

.github/workflows/validate.yml

@@ -128,6 +128,14 @@ jobs:
         env:
           SD_RUNTIME_API_TARGET: ${{ matrix.target }}
           SD_REQUIRE_RUNTIME: "1"
+      - name: Generate from FP8 weights
+        # Metal needs a Mac: GitHub's macOS runners list no Metal device.
+        if: runner.os != 'Windows' && matrix.smoke
+        run: python3 -m unittest discover -s tests -p test_fp8_runtime.py -v
+        env:
+          SD_FP8_TARGET: ${{ matrix.target }}
+          SD_REQUIRE_RUNTIME: "1"
+          GGML_VK_VISIBLE_DEVICES: "0"
       - uses: dart-lang/setup-dart@v1
         if: runner.os == 'Linux' && matrix.smoke
       - name: Verify cooperative host shutdown

CMakeLists.txt

@@ -76,6 +76,7 @@ function(sd_native_carry_patch patch)
 endfunction()
 
 sd_native_carry_patch(pickle_line_span.patch src/model_io/pickle_io.cpp)
+sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp)
 
 # Wrapper exports are compiled into the upstream library target, so they share
 # its SD_API export macro and reach upstream without patching it.

README.md

@@ -610,9 +610,10 @@ the macOS probe and builds the iOS device and simulator slices.
 Upstream 1017 has its
 [ABI audit, wrapper review, FP8 evidence and outstanding gates](docs/upstream_1017_qualification.md);
 the carry inventory is in the [upstream 982 record](docs/upstream_982_qualification.md).
-It carries two fixes the pin lacks: unbatched CFG in Metal builds (below) and
+It carries three fixes the pin lacks: unbatched CFG in Metal builds (below),
 the [`pickle_line_span.patch`](patches/README.md) for legacy checkpoints above
-2 GiB.
+2 GiB, and the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8
+weights as F16 on Metal instead of aborting.
 Upstream's `generate_image` now returns an owned result with per-image
 generation metadata; `sd_dart_generate_image` keeps returning an image array
 that `free_sd_images` frees and drops that metadata. Since upstream 982 a null

docs/upstream_1017_qualification.md

@@ -113,6 +113,11 @@ This is a pre-existing defect of every published Apple runtime, not caused or
 fixed by this upgrade, and the default Apple backend crashes rather than
 refusing. Vulkan was not run.
 
+`patches/fp8_metal_f16_load.patch` ([#29](https://github.com/leehack/stable-diffusion-native/issues/29))
+fixes it after this upgrade: under a debugger the abort is in
+`ggml_metal_op_mul_mat` during the SD2 v-prediction probe in `new_sd_ctx`, and
+Metal now loads FP8 weights as F16.
+
 ## Executed qualification
 
 macOS arm64 (Apple M4 Max), release build of this branch, dylib SHA-256

patches/README.md

@@ -13,3 +13,4 @@ an upstream issue or pull request to retire it.
 | Patch | Sources | Fixes | Test |
 | --- | --- | --- | --- |
 | `pickle_line_span.patch` | `src/model_io/pickle_io.cpp` | A legacy PyTorch checkpoint above 2 GiB is rejected: the pickle scanner narrows the rest of the file to `int` before looking for a newline ([#23](https://github.com/leehack/stable-diffusion-native/issues/23)). | `tests/test_pickle_span.py` |
+| `fp8_metal_f16_load.patch` | `src/model_manager.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). | `tests/test_fp8_runtime.py` |

patches/fp8_metal_f16_load.patch

@@ -0,0 +1,56 @@
+diff --git a/src/model_manager.cpp b/src/model_manager.cpp
+--- a/src/model_manager.cpp
++++ b/src/model_manager.cpp
+@@ -2,6 +2,7 @@
+
+ #include <algorithm>
+ #include <cstdint>
++#include <cstring>
+ #include <iterator>
+ #include <mutex>
+ #include <tuple>
+@@ -185,6 +186,36 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
+     return size;
+ }
+
++// ggml-metal reports FP8 MUL_MAT as supported, so Linear keeps FP8 weights, but
++// it has no FP8 matrix kernel and aborts the process encoding one. Weights a
++// Metal backend computes with are loaded as F16 instead, which holds every FP8
++// value exactly.
++static void load_fp8_params_as_f16_on_metal(const std::map<std::string, ggml_tensor*>& tensors,
++                                            ggml_backend_t compute_backend) {
++#ifndef SD_USE_UPSTREAM_GGML
++    ggml_backend_dev_t device = compute_backend != nullptr ? ggml_backend_get_device(compute_backend) : nullptr;
++    ggml_backend_reg_t reg    = device != nullptr ? ggml_backend_dev_backend_reg(device) : nullptr;
++    if (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "MTL") != 0) {
++        return;
++    }
++    for (const auto& pair : tensors) {
++        ggml_tensor* tensor = pair.second;
++        if (tensor == nullptr || (tensor->type != GGML_TYPE_F8_E4M3 && tensor->type != GGML_TYPE_F8_E5M2) ||
++            tensor->data != nullptr || tensor->buffer != nullptr || tensor->view_src != nullptr) {
++            continue;
++        }
++        tensor->type  = GGML_TYPE_F16;
++        tensor->nb[0] = ggml_type_size(GGML_TYPE_F16);
++        for (int i = 1; i < GGML_MAX_DIMS; ++i) {
++            tensor->nb[i] = tensor->nb[i - 1] * tensor->ne[i - 1];
++        }
++    }
++#else
++    (void)tensors;
++    (void)compute_backend;
++#endif
++}
++
+ void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits) {
+     if (compute_backend == nullptr) {
+         return;
+@@ -228,6 +259,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
+         LOG_ERROR("model manager tensor component is invalid");
+         return false;
+     }
++    load_fp8_params_as_f16_on_metal(tensors, compute_backend);
+     if (registered_tensor_size != nullptr) {
+         *registered_tensor_size += estimate_tensors_size(tensors);
+     }

tests/native/fp8_runtime_test.cpp

@@ -0,0 +1,91 @@
+// Loads a model and generates one image through the guarded wrapper exports,
+// as llamadart does, and writes its RGB bytes.
+//
+// Usage: fp8_runtime_test devices
+//        fp8_runtime_test generate <model> <backend> <out.rgb>
+//
+// The model is one that `exit_teardown_runtime_test make-model` wrote, or an
+// FP8 copy of it: it holds its TAESD decoder, so it needs no VAE.
+
+#include "sd_dart_wrapper.h"
+
+#include <cstdio>
+#include <cstring>
+#include <string>
+#include <vector>
+
+namespace {
+
+const int kImageSize = 64;
+
+int list_devices() {
+    std::vector<char> devices(sd_list_devices(nullptr, 0) + 1);
+    sd_list_devices(devices.data(), devices.size());
+    std::printf("%s", devices.data());
+    return 0;
+}
+
+int generate(const char* model, const char* backend, const char* output) {
+    sd_dart_log_enable();
+    sd_ctx_params_t context_params;
+    sd_ctx_params_init(&context_params);
+    context_params.model_path = model;
+    context_params.taesd_path = model;
+    context_params.backend    = backend;
+    sd_ctx_t* context         = sd_dart_new_sd_ctx(&context_params);
+    if (context == nullptr) {
+        std::vector<char> error(4096);
+        sd_dart_last_error(error.data(), error.size());
+        std::fprintf(stderr, "load failed: %s\n", error.data());
+        return 3;
+    }
+
+    sd_img_gen_params_t params;
+    sd_img_gen_params_init(&params);
+    params.prompt                         = "a lighthouse";
+    params.negative_prompt                = "";
+    params.width                          = kImageSize;
+    params.height                         = kImageSize;
+    params.seed                           = 7;
+    params.sample_params.sample_steps     = 2;
+    params.sample_params.guidance.txt_cfg = 1.0f;
+    sd_image_t* images                    = nullptr;
+    int count                             = 0;
+    if (!sd_dart_generate_image(context, &params, &images, &count)) {
+        std::vector<char> error(4096);
+        sd_dart_last_error(error.data(), error.size());
+        std::fprintf(stderr, "generation failed: %s\n", error.data());
+        sd_dart_exit_free(context);
+        return 4;
+    }
+    int status = 0;
+    if (images == nullptr || count != 1 || images[0].width != kImageSize ||
+        images[0].height != kImageSize || images[0].channel != 3) {
+        std::fprintf(stderr, "unexpected images\n");
+        status = 5;
+    } else {
+        const size_t bytes = static_cast<size_t>(kImageSize) * kImageSize * 3;
+        FILE* file         = std::fopen(output, "wb");
+        if (file == nullptr || std::fwrite(images[0].data, 1, bytes, file) != bytes ||
+            std::fclose(file) != 0) {
+            std::fprintf(stderr, "cannot write %s\n", output);
+            status = 6;
+        }
+    }
+    free_sd_images(images, count);
+    sd_dart_exit_free(context);
+    return status;
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    if (argc == 2 && std::strcmp(argv[1], "devices") == 0) {
+        return list_devices();
+    }
+    if (argc == 5 && std::strcmp(argv[1], "generate") == 0) {
+        return generate(argv[2], argv[3], argv[4]);
+    }
+    std::fprintf(stderr, "usage: %s devices | generate <model> <backend> <out.rgb>\n", argv[0]);
+    return 2;
+}