From a017f7a2325b1ea13b6f6ada5ca9373d4a140b65 Mon Sep 17 00:00:00 2001 From: Pete Thomason Date: Tue, 15 Sep 2026 11:29:08 -0500 Subject: [PATCH] fix: handle multi-dim activations in LoKr forward ggml_ext_lokr_forward's non-conv branch assumed a 2-D activation [q, batch] and took batch = ne[1]. Model blocks that run linear layers on 3-D activations [features, L, N] (e.g. Krea-2 attention with N > 1) hit the first split reshape with a mismatched element count, and the GGML_ASSERT inside ggml_reshape_3d aborted the process (SIGABRT) whenever a LoKr LoRA was active on such a layer. Fold every trailing dim into the batch (nelements / q) and reshape the result back to the activation's original trailing dims, so the out-diff adds element-wise onto the linear output instead of relying on a flat 2-D shape. The conv branch already batches on ne[3] and is unchanged. --- src/model/adapter/lora_ops.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/model/adapter/lora_ops.cpp b/src/model/adapter/lora_ops.cpp index 7dacfd164..c04dd5895 100644 --- a/src/model/adapter/lora_ops.cpp +++ b/src/model/adapter/lora_ops.cpp @@ -65,7 +65,11 @@ ggml_tensor* ggml_ext_lokr_forward( ggml_tensor* hb; if (!is_conv) { - int batch = (int)h->ne[1]; + // The activation may carry more than one trailing dim (e.g. + // [features, L, N] in batched attention); fold them all into + // the batch and restore the original layout on the way out — + // ne[1] alone under-counts and the split reshape below asserts. + int batch = (int)(ggml_nelements(h) / q_actual); int merge_batch_uq = batch; int merge_batch_vp = batch; @@ -117,8 +121,12 @@ ggml_tensor* ggml_ext_lokr_forward( hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch); } - ggml_tensor* hc = ggml_transpose(ctx, hc_t); - ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch); + ggml_tensor* hc = ggml_transpose(ctx, hc_t); + // Restore the activation's trailing dims so the diff adds + // element-wise onto the linear output instead of broadcasting + // a flattened batch across it. + ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), up * vp, + h->ne[1], h->ne[2], h->ne[3]); return ggml_ext_scale(ctx, out, scale); } else { int batch = (int)h->ne[3];