feat: add circular RoPE support for ideogram4 (#1627)

src/core/ggml_extend.hpp

@@ -1346,10 +1346,18 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
 
         if (mask_in != nullptr) {
-            mask_in = ggml_transpose(ctx, mask_in);
-        }
-
-        if (mask_in != nullptr) {
+            // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped
+            // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and,
+            // unlike the manual-attention path, does not broadcast the query dimension.
+            // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over
+            // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the
+            // kernel indexes it correctly. (A bare ggml_transpose here produced a
+            // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank
+            // output for masked flash attention.)
+            if (mask_in->ne[1] != L_q) {
+                mask_in = ggml_repeat(ctx, mask_in,
+                                      ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3]));
+            }
             mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
         }
 
@@ -2470,10 +2478,26 @@ protected:
             *effective_budget_out = effective_budget;
         }
 
+        // When streaming and the model dwarfs the budget, cap the planner at
+        // a quarter so it builds smaller merged segments and chunk-K can fit
+        // alongside. Without streaming the cap only adds dispatch overhead.
+        size_t planner_budget = effective_budget;
+        if (stream_layers_enabled) {
+            size_t total_params_bytes = 0;
+            for (const ggml_tensor* t : params_tensor_set_) {
+                if (t != nullptr) {
+                    total_params_bytes += ggml_nbytes(t);
+                }
+            }
+            if (total_params_bytes * 4 > effective_budget * 3) {
+                planner_budget = effective_budget / 4;
+            }
+        }
+
         *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
                                                      gf,
                                                      &graph_cut_plan_cache_,
-                                                     effective_budget,
+                                                     planner_budget,
                                                      params_tensor_set_,
                                                      get_desc().c_str());
         if (stream_layers_enabled) {
@@ -3311,6 +3335,7 @@ public:
         for (auto& pair : params) {
             ggml_tensor* param           = pair.second;
             tensors[prefix + pair.first] = pair.second;
+            ggml_set_name(param, (prefix + pair.first).c_str());
         }
     }
 

src/model/common/rope.hpp

@@ -253,7 +253,8 @@ namespace Rope {
                                                                  int bs,
                                                                  float theta,
                                                                  int head_dim,
-                                                                 const std::vector<int>& mrope_section) {
+                                                                 const std::vector<int>& mrope_section,
+                                                                 const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
         GGML_ASSERT(bs > 0);
         GGML_ASSERT(head_dim % 2 == 0);
         GGML_ASSERT(mrope_section.size() >= 3);
@@ -265,7 +266,11 @@ namespace Rope {
         std::vector<std::vector<std::vector<float>>> axis_embs;
         axis_embs.reserve(3);
         for (int axis = 0; axis < 3; ++axis) {
-            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta));
+            std::vector<int> axis_wrap;
+            if (axis < static_cast<int>(axis_wrap_dims.size())) {
+                axis_wrap = axis_wrap_dims[axis];
+            }
+            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
         }
 
         std::vector<std::vector<float>> emb = axis_embs[0];

src/model/diffusion/ideogram4.hpp

@@ -151,7 +151,9 @@ namespace Ideogram4 {
                                                           int context_len,
                                                           int head_dim,
                                                           int rope_theta,
-                                                          const std::vector<int>& mrope_section) {
+                                                          const std::vector<int>& mrope_section,
+                                                          bool circular_x = false,
+                                                          bool circular_y = false) {
         GGML_ASSERT(bs == 1);
         std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
                                             std::vector<float>(3, 0.f));
@@ -169,7 +171,29 @@ namespace Ideogram4 {
             }
         }
 
-        return Rope::embed_interleaved_mrope(ids, bs, static_cast<float>(rope_theta), head_dim, mrope_section);
+        std::vector<std::vector<int>> axis_wrap_dims(3);
+        if (circular_y || circular_x) {
+            size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
+            axis_wrap_dims[1].assign(total_len, 0);
+            axis_wrap_dims[2].assign(total_len, 0);
+            if (circular_y) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[1][idx] = grid_h;
+                }
+            }
+            if (circular_x) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[2][idx] = grid_w;
+                }
+            }
+        }
+
+        return Rope::embed_interleaved_mrope(ids,
+                                            bs,
+                                            static_cast<float>(rope_theta),
+                                            head_dim,
+                                            mrope_section,
+                                            axis_wrap_dims);
     }
 
     class Ideogram4Attention : public GGMLBlock {
@@ -480,13 +504,16 @@ namespace Ideogram4 {
             int64_t pos_len  = context_len + grid_h * grid_w;
             int64_t head_dim = config.emb_dim / config.num_heads;
 
+            auto runner_ctx = get_context();
             pe_vec  = gen_ideogram4_pe(static_cast<int>(grid_h),
                                        static_cast<int>(grid_w),
                                        static_cast<int>(x->ne[3]),
                                        static_cast<int>(context_len),
                                        static_cast<int>(head_dim),
                                        static_cast<int>(config.rope_theta),
-                                       config.mrope_section);
+                                       config.mrope_section,
+                                       runner_ctx.circular_x_enabled,
+                                       runner_ctx.circular_y_enabled);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
 
@@ -497,7 +524,6 @@ namespace Ideogram4 {
             auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]);
             set_backend_tensor_data(indicator, image_indicator_vec.data());
 
-            auto runner_ctx  = get_context();
             ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator);
             ggml_build_forward_expand(gf, out);
             return gf;

src/runtime/guidance.cpp

@@ -173,8 +173,9 @@ namespace sd::guidance {
         }
 
         float diff_norm = 0.0f;
+        const int standard_res = 2 * 1024 / 8; // Use SDXL as the standard resolution (1024x1024, 8x8 patches, 4=2x2 channels)
         if (params_.norm_threshold > 0.0f) {
-            diff_norm = std::sqrt((deltas * deltas).sum());
+            diff_norm = std::sqrt((deltas * deltas).sum()) * standard_res / std::sqrt(static_cast<float>(deltas.numel()));
         }
 
         float apg_scale_factor = 1.0f;

src/stable-diffusion.cpp

@@ -575,15 +575,6 @@ public:
                     }
                 }
                 if (is_chroma) {
-                    if ((sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) && sd_ctx_params->chroma_use_dit_mask) {
-                        LOG_WARN(
-                            "!!!It looks like you are using Chroma with flash attention. "
-                            "This is currently unsupported. "
-                            "If you find that the generated images are broken, "
-                            "try either disabling flash attention or specifying "
-                            "--chroma-disable-dit-mask as a workaround.");
-                    }
-
                     cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
                                                                         params_backend_for(SDBackendModule::TE),
                                                                         tensor_storage_map,