fix: avoid f16 overflow in Z-Image quantized matmuls on CUDA (#2074)

Co-authored-by: leejet <leejet714@gmail.com>

src/model/common/ggml_block.hpp

@@ -205,6 +205,11 @@ public:
         force_prec_f32 = force_prec_f32_;
     }
 
+    bool has_quantized_weight() const {
+        auto it = params.find("weight");
+        return it != params.end() && it->second != nullptr && ggml_is_quantized(it->second->type);
+    }
+
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;

src/model/diffusion/z_image.hpp

@@ -168,13 +168,16 @@ namespace ZImage {
             int64_t N       = x->ne[2];
             auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
 
+            if (out_proj->has_quantized_weight()) {
+                out_proj->set_scale(1.f / 16.f);
+            }
+
             if (split_qkv) {
                 auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
                 auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
 
                 if (sd_backend_is(ctx->backend, "ROCm")) {
-                    out_proj->set_scale(1.f / 16.f);
                     out_proj->set_force_prec_f32(true);
                     q_proj->set_force_prec_f32(true);
                     k_proj->set_force_prec_f32(true);
@@ -197,7 +200,6 @@ namespace ZImage {
             auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
             if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_scale(1.f / 16.f);
                 out_proj->set_force_prec_f32(true);
                 qkv_proj->set_force_prec_f32(true);
             }