@@ -16654,46 +16654,68 @@ struct llm_build_lfm2 : public llm_graph_context {
1665416654 ggml_tensor * cur,
1665516655 llm_graph_input_rs * inp_recr,
1665616656 int il) {
16657- const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx)->get_recr();
16657+ const auto * mctx_cur = static_cast<const llama_memory_hybrid_context *>(mctx)->get_recr();
16658+ const uint32_t kv_head = mctx_cur->get_head();
16659+ const int64_t n_seq_tokens = ubatch.n_seq_tokens;
16660+ const int64_t n_seqs = ubatch.n_seqs;
16661+ GGML_ASSERT(n_seqs != 0);
16662+ GGML_ASSERT(ubatch.equal_seqs);
16663+ GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
16664+
16665+ GGML_ASSERT(hparams.n_shortconv_l_cache > 1);
16666+ const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
16667+
16668+ // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
16669+ cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
1665816670
1665916671 auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
1666016672 cb(bcx, "model.layers.{}.conv.in_proj", il);
1666116673
1666216674 constexpr auto n_chunks = 3;
1666316675 GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
1666416676 auto const chunk_size = bcx->ne[0] / n_chunks;
16665- auto * b = ggml_view_2d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->nb[1], 0 * chunk_size * ggml_element_size(bcx));
16666- auto * c = ggml_view_2d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->nb[1], 1 * chunk_size * ggml_element_size(bcx));
16667- auto * x = ggml_view_2d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->nb[1], 2 * chunk_size * ggml_element_size(bcx));
16677+ auto * b = ggml_view_3d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx-> nb[1], bcx->nb[2], 0* chunk_size* ggml_element_size(bcx));
16678+ auto * c = ggml_view_3d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx-> nb[1], bcx->nb[2], 1* chunk_size* ggml_element_size(bcx));
16679+ auto * x = ggml_view_3d (ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx-> nb[1], bcx->nb[2], 2* chunk_size* ggml_element_size(bcx));
1666816680
1666916681 auto * bx = ggml_transpose(ctx0, ggml_mul(ctx0, b, x));
1667016682
16671- // read conv state directly, with build_rs generation is slower
16672- ggml_tensor * conv_state = mctx_cur->get_r_l(il);
16673- const int64_t n_seqs = ubatch.n_seqs;
16674- ggml_tensor * conv = build_rs(inp_recr, gf, conv_state, hparams.n_embd_r(), n_seqs);
16675- conv = ggml_reshape_3d(ctx0, conv_state, hparams.n_shortconv_l_cache - 1, hparams.n_embd, n_seqs);
16683+ // read conv state
16684+ auto * conv_state = mctx_cur->get_r_l(il);
16685+ auto * conv_rs = build_rs(inp_recr, gf, conv_state, hparams.n_embd_r(), n_seqs);
16686+ auto * conv = ggml_reshape_3d(ctx0, conv_rs, d_conv, hparams.n_embd, n_seqs);
1667616687
1667716688 bx = ggml_concat(ctx0, conv, bx, 0);
1667816689 GGML_ASSERT(bx->ne[0] > conv->ne[0]);
1667916690
16680- auto * new_conv = ggml_view_2d(ctx0, bx, conv->ne[0], bx->ne[1], bx->nb[1], (bx->ne[0] - conv->ne[0]) * ggml_element_size(bx));
16691+ // last d_conv columns is a new conv state
16692+ auto * new_conv = ggml_view_3d(ctx0, bx, conv->ne[0], bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2], (bx->ne[0] - conv->ne[0])*ggml_element_size(bx));
1668116693 GGML_ASSERT(ggml_are_same_shape(conv, new_conv));
1668216694
16683- // write conv state
16684- ggml_build_forward_expand(gf, ggml_cpy(ctx0, new_conv, conv_state));
16695+ // write new conv conv state
16696+ ggml_build_forward_expand(
16697+ gf,
16698+ ggml_cpy(
16699+ ctx0,
16700+ new_conv,
16701+ ggml_view_1d(
16702+ ctx0,
16703+ conv_state,
16704+ ggml_nelements(new_conv),
16705+ kv_head*d_conv*n_embd*ggml_element_size(new_conv)
16706+ )
16707+ )
16708+ );
1668516709
1668616710 auto * conv_kernel = model.layers[il].shortconv.conv;
16687- GGML_ASSERT(hparams.n_shortconv_l_cache > 0);
16688-
16689- // construct ssm_conv op
16690- ggml_tensor * conv_out = ggml_ssm_conv(ctx0, bx, conv_kernel);
16711+ auto * conv_out = ggml_ssm_conv(ctx0, bx, conv_kernel);
1669116712 cb(conv_out, "model.layers.{}.conv.conv", il);
1669216713
1669316714 auto * y = ggml_mul(ctx0, c, conv_out);
16694-
1669516715 y = build_lora_mm(model.layers[il].shortconv.out_proj, y);
1669616716 cb(y, "model.layers.{}.conv.out_proj", il);
16717+ // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
16718+ y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
1669716719
1669816720 return y;
1669916721 }
0 commit comments