reshaping so mm on 2d instead of 3d tensor for lfm sc (#30831)

This commit is contained in:
Anthony Maltsev authored and GitHub committed 2026-10-31 20:34:57 +01:00
1 parent f3586e66bc
commit ddf48f85c9
1 file changed
+4 -5
+4 -5
View File
@@ -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;
}