From ddf48f85c933c1f438212a02eb21e7b001d2d37d Mon Sep 17 00:00:00 2001
From: Anthony Maltsev <77708730+anthonymaltsev@users.noreply.github.com>
Date: Sat, 31 Oct 2026 12:34:57 -0700
Subject: [PATCH] reshaping so mm on 2d instead of 3d tensor for lfm sc
 (#30831)

---
 src/models/lfm-base.cpp | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/models/lfm-base.cpp b/src/models/lfm-base.cpp
index f5ff7d178..0dabf114b 100644
--- a/src/models/lfm-base.cpp
+++ b/src/models/lfm-base.cpp
@@ -22,11 +22,10 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
     const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
     const int64_t  n_embd = hparams.n_embd;
 
-    // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
-    cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
-
     auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
     cb(bcx, "model.layers.{}.conv.in_proj", il);
+    // {3*n_embd, n_tokens} => {3*n_embd, n_seq_tokens, n_seqs}
+    bcx = ggml_reshape_3d(ctx0, bcx, bcx->ne[0], n_seq_tokens, n_seqs);
 
     constexpr auto n_chunks = 3;
     GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
@@ -76,10 +75,10 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
     cb(conv_out, "model.layers.{}.conv.conv", il);
 
     auto * y = ggml_mul(ctx0, c, conv_out);
-    y        = build_lora_mm(model.layers[il].shortconv.out_proj, y);
-    cb(y, "model.layers.{}.conv.out_proj", il);
     // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
     y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
+    y = build_lora_mm(model.layers[il].shortconv.out_proj, y);
+    cb(y, "model.layers.{}.conv.out_proj", il);
 
     return y;
 }