#include "models.h" // streaming encoder state: base position, conv left context, short-conv taps and attention window // a zero state gives the same result as encoding the whole clip std::vector list_lfm3a_state_slots(const clip_hparams & hparams, const clip_model & model) { std::vector slots; const int64_t n_mel = hparams.n_mel_bins; const int64_t n_ch = model.pre_encode_conv_X_w[0]->ne[3]; const int64_t prefix = hparams.attn_window_size - 1; slots.push_back({"pos", 1, 1}); slots.push_back({"ds0", 1, n_mel}); slots.push_back({"ds1", 1, (n_mel / 2) * n_ch}); slots.push_back({"ds2", 1, (n_mel / 4) * n_ch}); for (size_t il = 0; il < model.layers.size(); il++) { const auto & layer = model.layers[il]; if (layer.conv_dw_w) { slots.push_back({"conv_" + std::to_string(il), layer.conv_dw_w->ne[0] - 1, hparams.n_embd}); } else { slots.push_back({"k_" + std::to_string(il), layer.k_w->ne[1], prefix}); slots.push_back({"v_" + std::to_string(il), layer.v_w->ne[1], prefix}); } } return slots; } // LFM3-Audio encoder: conv downsampler (8x in time) -> causal LFM2 stack -> MLP projector ggml_cgraph * clip_graph_lfm3a::build() { const norm_type norm_t = NORM_TYPE_RMS; const float rope_theta = hparams.rope_theta; const int64_t prefix = hparams.attn_window_size - 1; // the graph always streams, a zero state encodes a clip from its start std::map state_in; std::vector> state_out; for (const auto & slot : list_lfm3a_state_slots(hparams, model)) { ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); ggml_set_name(t, ("state_in_" + slot.name).c_str()); ggml_set_input(t); state_in[slot.name] = t; } // stride-2 conv over [time, freq, ch]; the left column comes from the state and the right edge is zero-padded auto conv_down = [&](ggml_tensor * x, const std::string & name, ggml_tensor * w, ggml_tensor * b, bool depthwise) { ggml_tensor * left = ggml_reshape_4d(ctx0, state_in.at(name), 1, x->ne[1], x->ne[2], 1); x = ggml_concat(ctx0, left, x, 0); ggml_tensor * last = ggml_view_4d(ctx0, x, 1, x->ne[1], x->ne[2], 1, x->nb[1], x->nb[2], x->nb[3], (x->ne[0] - 1) * x->nb[0]); state_out.push_back({name, ggml_reshape_2d(ctx0, ggml_cont(ctx0, last), 1, x->ne[1] * x->ne[2])}); x = ggml_pad_ext(ctx0, x, 0, 1, 0, 0, 0, 0, 0, 0); x = depthwise ? ggml_conv_2d_dw_direct(ctx0, w, x, 2, 2, 0, 1, 1, 1) : ggml_conv_2d(ctx0, w, x, 2, 2, 0, 1, 1, 1); return ggml_add(ctx0, x, b); }; // inp: [n_frames, n_mel, 1] ggml_tensor * cur = build_inp_raw(1); // downsampler, same layout as the lfm2a pre-encoder but with SiLU cur = conv_down(cur, "ds0", model.pre_encode_conv_X_w[0], model.pre_encode_conv_X_b[0], false); cur = ggml_silu_inplace(ctx0, cur); cur = conv_down(cur, "ds1", model.pre_encode_conv_X_w[2], model.pre_encode_conv_X_b[2], true); cur = ggml_conv_2d_direct(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1); cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]); cur = ggml_silu_inplace(ctx0, cur); cur = conv_down(cur, "ds2", model.pre_encode_conv_X_w[5], model.pre_encode_conv_X_b[5], true); cur = ggml_conv_2d_direct(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1); cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]); cur = ggml_silu_inplace(ctx0, cur); cb(cur, "downsample.conv", -1); // [n_pos, n_mel/8, n_ch] -> [n_ch * n_mel/8, n_pos] const int n_pos = cur->ne[0]; cur = ggml_reshape_2d(ctx0, cur, n_pos, cur->ne[1] * cur->ne[2]); cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); cur = build_mm(model.pre_encode_out_w, cur); cur = ggml_add(ctx0, cur, model.pre_encode_out_b); cb(cur, "downsample.out", -1); ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); ggml_set_name(positions, "positions"); ggml_set_input(positions); state_out.push_back({"pos", ggml_scale_bias(ctx0, state_in.at("pos"), 1.0f, (float) n_pos)}); // causal mask with a sliding window, shared by all attention layers ggml_tensor * kq_mask = build_inp_attn_mask(prefix + n_pos, n_pos); ggml_tensor * inpL = cur; for (int il = 0; il < n_layer; il++) { const auto & layer = model.layers[il]; cur = build_norm(inpL, layer.ln_1_w, nullptr, norm_t, eps, il); cb(cur, "operator_norm", il); if (layer.conv_dw_w) { // short conv: in_proj -> b * x -> causal depthwise conv -> c * conv -> out_proj ggml_tensor * bcx = build_mm(layer.conv_pw1_w, cur); cb(bcx, "conv_in_proj", il); ggml_tensor * b = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], 0); ggml_tensor * c = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], ggml_row_size(bcx->type, n_embd)); ggml_tensor * x = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], ggml_row_size(bcx->type, 2 * n_embd)); ggml_tensor * bx = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mul(ctx0, b, x))); const int d_conv = layer.conv_dw_w->ne[0] - 1; const std::string name = "conv_" + std::to_string(il); bx = ggml_concat(ctx0, state_in.at(name), bx, 0); state_out.push_back({name, ggml_cont(ctx0, ggml_view_2d(ctx0, bx, d_conv, n_embd, bx->nb[1], (size_t) n_pos * bx->nb[0]))}); bx = ggml_reshape_3d(ctx0, bx, bx->ne[0], n_embd, 1); cur = ggml_ssm_conv(ctx0, bx, layer.conv_dw_w); cur = ggml_reshape_2d(ctx0, cur, n_embd, n_pos); cb(cur, "conv_conv", il); cur = ggml_mul(ctx0, c, cur); cur = build_mm(layer.conv_pw2_w, cur); cb(cur, "conv_out_proj", il); } else { ggml_tensor * Qcur = build_mm(layer.q_w, cur); ggml_tensor * Kcur = build_mm(layer.k_w, cur); ggml_tensor * Vcur = build_mm(layer.v_w, cur); Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head_kv, n_pos); Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head_kv, n_pos); Qcur = build_norm(Qcur, layer.q_norm, nullptr, norm_t, eps, il); Kcur = build_norm(Kcur, layer.k_norm, nullptr, norm_t, eps, il); Qcur = ggml_rope_ext(ctx0, Qcur, positions, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); Kcur = ggml_rope_ext(ctx0, Kcur, positions, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); // the cached window goes in front, this chunk's tail is kept for the next call const std::string k_name = "k_" + std::to_string(il); const std::string v_name = "v_" + std::to_string(il); const int64_t n_embd_kv = d_head * n_head_kv; ggml_tensor * K = ggml_concat(ctx0, state_in.at(k_name), ggml_reshape_2d(ctx0, Kcur, n_embd_kv, n_pos), 1); ggml_tensor * V = ggml_concat(ctx0, state_in.at(v_name), ggml_reshape_2d(ctx0, Vcur, n_embd_kv, n_pos), 1); state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, K, n_embd_kv, prefix, K->nb[1], (size_t) n_pos * K->nb[1]))}); state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, V, n_embd_kv, prefix, V->nb[1], (size_t) n_pos * V->nb[1]))}); Kcur = ggml_reshape_3d(ctx0, K, d_head, n_head_kv, prefix + n_pos); Vcur = ggml_reshape_3d(ctx0, V, d_head, n_head_kv, prefix + n_pos); cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il); cb(cur, "attn_out", il); } cur = ggml_add(ctx0, cur, inpL); inpL = cur; cur = build_norm(cur, layer.ln_2_w, nullptr, norm_t, eps, il); cb(cur, "ffn_norm", il); cur = build_ffn(cur, layer.ff_up_w, nullptr, layer.ff_gate_w, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il); cb(cur, "ffn_out", il); cur = ggml_add(ctx0, inpL, cur); inpL = cur; } cur = build_norm(inpL, model.post_ln_w, nullptr, norm_t, eps, -1); cb(cur, "embedding_norm", -1); // projector: LayerNorm -> Linear -> GELU -> Linear cur = build_norm(cur, model.mm_0_w, model.mm_0_b, NORM_TYPE_NORMAL, 1e-5, -1); cur = build_ffn(cur, model.mm_1_w, model.mm_1_b, nullptr, nullptr, model.mm_3_w, model.mm_3_b, FFN_GELU_ERF, -1); cb(cur, "projected", -1); // the state outputs go first so that the embeddings stay the last node of the graph for (const auto & s : state_out) { ggml_set_name(s.second, ("state_out_" + s.first).c_str()); ggml_set_output(s.second); ggml_build_forward_expand(gf, s.second); } ggml_build_forward_expand(gf, cur); return gf; }