#include "conformer.hpp"
#include "relpos_attention.hpp"
#include "ggml_graph.hpp"
#include "backend.hpp"
#include "ggml.h"
#include <cassert>
#include <cmath>
#include <cstring>
#include <string>
#include <vector>

namespace pk {

// Weights from the loader are referenced DIRECTLY as graph leaves via the shared
// pk::clone_weight / pk::clone_weight_opt (backend.cpp; zero-copy via the
// loader's CPU backend buffer). Quantization-allowlist linears may be f16/q8_0
// (ggml_mul_mat dequantizes src0 on the fly). Optional weights (conv biases,
// absent in bias=False checkpoints) use clone_weight_opt. std::string-name
// overloads keep the call sites (pre + suffix) unchanged.
static ggml_tensor* clone_weight(ggml_context* ctx, const ModelLoader& ml,
                                 const std::string& name) {
    return pk::clone_weight(ctx, ml, name.c_str());
}
static ggml_tensor* clone_weight_opt(ggml_context* ctx, const ModelLoader& ml,
                                     const std::string& name) {
    return pk::clone_weight_opt(ctx, ml, name.c_str());
}

// Build the ConformerConvolution sub-graph (everything AFTER norm_conv) on the
// conv input `c` (= norm_conv(residual)), ne [D, T, B] (channels fastest, batch
// on ne2). Returns the conv output tensor, ne [D, T, B]. See header for the NeMo
// `ConformerConvolution.forward` mapping. `valid_len[b]` is the per-item valid
// time length; the pre-depthwise pad mask is built per item so trailing-pad
// frames cannot leak through the time-mixing depthwise conv. All other ops are
// per-frame (1x1 convs == matmul, GLU, norm, silu) and broadcast over ne2.
// `pool` keeps host-side mask/scale/shift buffers alive until the enclosing
// Backend::compute finishes.
static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml,
                                      const std::string& pre, ggml_tensor* c,
                                      int D, int T, int K, int B,
                                      const std::vector<int>& valid_len,
                                      const std::string& conv_norm_type,
                                      bool conv_causal, GraphInputPool& pool) {
    const float ln_eps = 1e-5f;
    const float bn_eps = 1e-5f;
    const int pad = (K - 1) / 2;  // symmetric padding (offline model)

    // -- pointwise_conv1 (Conv1d d->2d, k=1): 1x1 conv == linear over channels.
    ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight");
    pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d]
    ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias");
    ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2d, T, B]
    if (pw1b) y = ggml_add(ctx, y, pw1b);

    // -- GLU over channel dim (NeMo F.glu(x, dim=1)). y is [2D, T, B]; each half
    // is D wide along ne0, stepping T via nb[1] and B via nb[2].
    ggml_tensor* a = ggml_view_3d(ctx, y, D, T, B, y->nb[1], y->nb[2], 0);
    ggml_tensor* b = ggml_view_3d(ctx, y, D, T, B, y->nb[1], y->nb[2],
                                  (size_t)D * y->nb[0]);
    ggml_tensor* glu = ggml_mul(ctx, ggml_cont(ctx, a),
                                ggml_sigmoid(ctx, ggml_cont(ctx, b))); // [d, T, B]

    // -- pad_mask: zero padded time positions before depthwise conv, per item.
    // The depthwise conv is the only time-mixing op here; without this mask the
    // trailing-pad region (no longer zero after pointwise bias/GLU) would leak
    // into the last valid output frame of a batched item. Mask is [1, T, B] and
    // broadcasts over ne0 (D). Emit only if some item has valid_len < T.
    bool need_mask = false;
    for (int bi = 0; bi < B; ++bi) if (valid_len[bi] < T) { need_mask = true; break; }
    if (need_mask) {
        std::vector<float>& md = pool.alloc_f32((size_t)T * B);
        for (int bi = 0; bi < B; ++bi)
            for (int t = 0; t < T; ++t)
                md[(size_t)bi * T + t] = (t < valid_len[bi]) ? 1.0f : 0.0f;
        int64_t tm_ne[3] = {1, T, B};
        ggml_tensor* tmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, tm_ne,
                                 md.data(), md.size() * sizeof(float));
        glu = ggml_mul(ctx, glu, tmask);
    }

    // -- depthwise_conv (Conv1d d->d, k=K, groups=d). F32 im2col throughout.
    // Transpose ne0<->ne1 (keep batch on ne2): [D,T,B] -> [T,C,B].
    ggml_tensor* glu_tcb = ggml_cont(ctx, ggml_permute(ctx, glu, 1, 0, 2, 3)); // [T, C, B]
    ggml_tensor* dww = clone_weight(ctx, ml, pre + "conv.depthwise_conv.weight"); // [K,1,C]
    ggml_tensor* dw = nullptr;
    {
        // ggml's 1D im2col asserts b->ne[3] == 1, so the depthwise conv cannot
        // take a batch dim. Run it per item (each slice has ne[3]==1, exactly the
        // B=1 path) and reassemble along the batch axis. For B==1 the loop runs
        // once and skips the concat, so the result is byte-identical to before.
        const int Tn = (int)glu_tcb->ne[0];
        const int Cn = (int)glu_tcb->ne[1];
        for (int bi = 0; bi < B; ++bi) {
            // item bi's [T, C, 1] slice out of glu_tcb [T, C, B]; cont for im2col.
            ggml_tensor* item = ggml_view_3d(ctx, glu_tcb, Tn, Cn, 1,
                                             glu_tcb->nb[1], glu_tcb->nb[2],
                                             (size_t)bi * glu_tcb->nb[2]); // [T,C,1]
            item = ggml_cont(ctx, item);
            // im2col layout: [W=T, H=1, C, N=1].
            ggml_tensor* nb = ggml_reshape_4d(ctx, item, Tn, 1, Cn, 1); // [T,1,C,1]
            ggml_tensor* ic;
            if (conv_causal) {
                ggml_tensor* nbp = ggml_pad_ext(ctx, nb, /*lp0*/K - 1, /*rp0*/0,
                                                0, 0, 0, 0, 0, 0); // [T+K-1,1,C,1]
                ic = ggml_im2col(ctx, dww, nbp, /*s0*/1, /*s1*/0,
                                 /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/0,
                                 /*is_2D*/false, GGML_TYPE_F32);
            } else {
                ic = ggml_im2col(ctx, dww, nb, /*s0*/1, /*s1*/0,
                                 /*p0*/pad, /*p1*/0, /*d0*/1, /*d1*/0,
                                 /*is_2D*/false, GGML_TYPE_F32);
            }
            // mul_mat result r2 ne = [OW=T, 1, C, 1]; drop the unit ne1 -> [T, C, 1].
            ggml_tensor* r2 = ggml_mul_mat(ctx, ic, dww);
            ggml_tensor* item_dw = ggml_reshape_3d(ctx, r2, r2->ne[0], r2->ne[2], 1); // [T,C,1]
            // reassemble along ne2 (batch). For B==1 there is no concat.
            dw = (dw == nullptr) ? item_dw : ggml_concat(ctx, dw, item_dw, /*dim*/2);
        }
        // dw is [OW=T, C, B].
    }
    ggml_tensor* dwb = clone_weight_opt(ctx, ml, pre + "conv.depthwise_conv.bias"); // [C]
    // transpose ne0<->ne1 (keep batch on ne2): [T,C,B] -> [C,T,B].
    ggml_tensor* dwt = ggml_cont(ctx, ggml_permute(ctx, dw, 1, 0, 2, 3)); // [C, T, B]
    if (dwb) dwt = ggml_add(ctx, dwt, dwb);               // broadcast [C] over T,B

    // -- norm (between depthwise conv and SiLU).
    ggml_tensor* normed;
    if (conv_norm_type == "layer_norm") {
        ggml_tensor* g = clone_weight(ctx, ml, pre + "conv.batch_norm.weight"); // [C]
        ggml_tensor* bb = clone_weight(ctx, ml, pre + "conv.batch_norm.bias");  // [C]
        normed = ggml_norm(ctx, dwt, ln_eps);             // normalize over ne0=C
        normed = ggml_mul(ctx, normed, g);                // *gamma  (broadcast [C] over T,B)
        normed = ggml_add(ctx, normed, bb);               // +beta
    } else {
        // batch_norm (inference): fold into per-channel scale/shift constants:
        //   scale = g / sqrt(var+eps); shift = b - mean*scale.  Computed host-side.
        std::vector<float>& sc = pool.alloc_f32(D);
        std::vector<float>& sh = pool.alloc_f32(D);
        std::vector<float> g, bb, m, var;
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.weight").c_str(), g);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.bias").c_str(), bb);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_mean").c_str(), m);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_var").c_str(), var);
        for (int cc = 0; cc < D; ++cc) {
            sc[cc] = g[cc] / std::sqrt(var[cc] + bn_eps);
            sh[cc] = bb[cc] - m[cc] * sc[cc];
        }
        int64_t d_ne[1] = {D};
        ggml_tensor* scale = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne,
                                 sc.data(), sc.size() * sizeof(float));
        ggml_tensor* shift = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne,
                                 sh.data(), sh.size() * sizeof(float));
        normed = ggml_add(ctx, ggml_mul(ctx, dwt, scale), shift); // [C, T, B]
    }

    // -- SiLU (Swish), then pointwise_conv2 (Conv1d d->d, k=1).
    normed = ggml_silu(ctx, normed);
    ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight");
    pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d]
    ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias");
    ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T, B]
    if (pw2b) cout = ggml_add(ctx, cout, pw2b);
    return cout; // [D, T, B]; this is layers[i].conv output
}

// Scalar (B=1) overload: the verbatim v1 2-D ConformerConvolution sub-graph on
// conv input `c`, ne [D, T] (channels fastest). Used by the single-clip /
// forward path (build_graph below, forward_with_conv localization, and
// conv_module_forward) so B=1 runs the lean 2-D graph and is bit-exact with v1.
// Distinguished from the batched overload by signature: this takes `int
// valid_len`, the batched one takes `int B, const std::vector<int>&`.
static ggml_tensor* build_conv_module(ggml_context* ctx, const ModelLoader& ml,
                                      const std::string& pre, ggml_tensor* c,
                                      int D, int T, int K, int valid_len,
                                      const std::string& conv_norm_type,
                                      bool conv_causal, GraphInputPool& pool) {
    const float ln_eps = 1e-5f;
    const float bn_eps = 1e-5f;
    const int pad = (K - 1) / 2;  // symmetric padding (offline model)

    // -- pointwise_conv1 (Conv1d d->2d, k=1): 1x1 conv == linear over channels.
    ggml_tensor* pw1w = clone_weight(ctx, ml, pre + "conv.pointwise_conv1.weight");
    pw1w = ggml_reshape_2d(ctx, pw1w, D, 2 * D); // [in=d, out=2d]
    ggml_tensor* pw1b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv1.bias");
    ggml_tensor* y = ggml_mul_mat(ctx, pw1w, c); // [2d, T]
    if (pw1b) y = ggml_add(ctx, y, pw1b);

    // -- GLU over channel dim (NeMo F.glu(x, dim=1)).
    ggml_tensor* a = ggml_view_2d(ctx, y, D, T, y->nb[1], 0);
    ggml_tensor* b = ggml_view_2d(ctx, y, D, T, y->nb[1], (size_t)D * y->nb[0]);
    ggml_tensor* glu = ggml_mul(ctx, ggml_cont(ctx, a),
                                ggml_sigmoid(ctx, ggml_cont(ctx, b))); // [d, T]

    // -- pad_mask: zero padded time positions before depthwise conv.
    if (valid_len < T) {
        std::vector<float>& md = pool.alloc_f32(T);
        for (int t = 0; t < T; ++t) md[t] = (t < valid_len) ? 1.0f : 0.0f;
        int64_t tm_ne[2] = {1, T};
        ggml_tensor* tmask = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, tm_ne,
                                 md.data(), md.size() * sizeof(float));
        glu = ggml_mul(ctx, glu, tmask);
    }

    // -- depthwise_conv (Conv1d d->d, k=K, groups=d). F32 im2col throughout.
    ggml_tensor* glu_tc = ggml_cont(ctx, ggml_transpose(ctx, glu)); // [T, C]
    ggml_tensor* dww = clone_weight(ctx, ml, pre + "conv.depthwise_conv.weight"); // [K,1,C]
    ggml_tensor* dw;
    {
        ggml_tensor* nb = ggml_reshape_4d(ctx, glu_tc,
                              glu_tc->ne[0], 1, glu_tc->ne[1], 1); // [T,1,C,1]
        ggml_tensor* ic;
        if (conv_causal) {
            ggml_tensor* nbp = ggml_pad_ext(ctx, nb, /*lp0*/K - 1, /*rp0*/0,
                                            0, 0, 0, 0, 0, 0); // [T+K-1,1,C,1]
            ic = ggml_im2col(ctx, dww, nbp, /*s0*/1, /*s1*/0,
                             /*p0*/0, /*p1*/0, /*d0*/1, /*d1*/0,
                             /*is_2D*/false, GGML_TYPE_F32);
        } else {
            ic = ggml_im2col(ctx, dww, nb, /*s0*/1, /*s1*/0,
                             /*p0*/pad, /*p1*/0, /*d0*/1, /*d1*/0,
                             /*is_2D*/false, GGML_TYPE_F32);
        }
        ggml_tensor* r2 = ggml_mul_mat(ctx, ic, dww);
        dw = ggml_reshape_3d(ctx, r2, r2->ne[0], r2->ne[2], 1); // [OW=T, C, 1]
    }
    dw = ggml_reshape_2d(ctx, dw, T, D);                  // [T, C]
    ggml_tensor* dwb = clone_weight_opt(ctx, ml, pre + "conv.depthwise_conv.bias"); // [C]
    ggml_tensor* dwt = ggml_cont(ctx, ggml_transpose(ctx, dw)); // [C, T]
    if (dwb) dwt = ggml_add(ctx, dwt, dwb);               // broadcast [C] over T

    // -- norm (between depthwise conv and SiLU).
    ggml_tensor* normed;
    if (conv_norm_type == "layer_norm") {
        ggml_tensor* g = clone_weight(ctx, ml, pre + "conv.batch_norm.weight"); // [C]
        ggml_tensor* bb = clone_weight(ctx, ml, pre + "conv.batch_norm.bias");  // [C]
        normed = ggml_norm(ctx, dwt, ln_eps);             // normalize over ne0=C
        normed = ggml_mul(ctx, normed, g);                // *gamma  (broadcast [C] over T)
        normed = ggml_add(ctx, normed, bb);               // +beta
    } else {
        // batch_norm (inference): fold into per-channel scale/shift constants:
        //   scale = g / sqrt(var+eps); shift = b - mean*scale.  Computed host-side.
        std::vector<float>& sc = pool.alloc_f32(D);
        std::vector<float>& sh = pool.alloc_f32(D);
        std::vector<float> g, bb, m, var;
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.weight").c_str(), g);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.bias").c_str(), bb);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_mean").c_str(), m);
        pk::weight_to_host_f32(ml, (pre + "conv.batch_norm.running_var").c_str(), var);
        for (int cc = 0; cc < D; ++cc) {
            sc[cc] = g[cc] / std::sqrt(var[cc] + bn_eps);
            sh[cc] = bb[cc] - m[cc] * sc[cc];
        }
        int64_t d_ne[1] = {D};
        ggml_tensor* scale = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne,
                                 sc.data(), sc.size() * sizeof(float));
        ggml_tensor* shift = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, d_ne,
                                 sh.data(), sh.size() * sizeof(float));
        normed = ggml_add(ctx, ggml_mul(ctx, dwt, scale), shift); // [C, T]
    }

    // -- SiLU (Swish), then pointwise_conv2 (Conv1d d->d, k=1).
    normed = ggml_silu(ctx, normed);
    ggml_tensor* pw2w = clone_weight(ctx, ml, pre + "conv.pointwise_conv2.weight");
    pw2w = ggml_reshape_2d(ctx, pw2w, D, D); // [in=d, out=d]
    ggml_tensor* pw2b = clone_weight_opt(ctx, ml, pre + "conv.pointwise_conv2.bias");
    ggml_tensor* cout = ggml_mul_mat(ctx, pw2w, normed); // [d, T]
    if (pw2b) cout = ggml_add(ctx, cout, pw2b);
    return cout; // [D, T] -> row-major [T, D]; this is layers[i].conv output
}

ConformerLayer::ConformerLayer(const ModelLoader& ml, int layer_idx)
    : ml_(ml), layer_idx_(layer_idx) {
    d_model_     = (int)ml.config().d_model;
    n_heads_     = (int)ml.config().n_heads;
    ff_dim_      = (int)ml.config().ff_dim;
    conv_kernel_ = (int)ml.config().conv_kernel;
    conv_norm_type_ = ml.config().conv_norm_type;
    conv_causal_    = ml.config().conv_causal;
    assert((conv_norm_type_ == "batch_norm" || conv_norm_type_ == "layer_norm") &&
           "ConformerLayer supports conv_norm_type in {batch_norm, layer_norm}");
    assert(n_heads_ > 0 && d_model_ % n_heads_ == 0);
    assert((conv_kernel_ - 1) % 2 == 0 && "depthwise kernel must be odd");
}

ggml_tensor* ConformerLayer::build_graph_batched(ggml_context* ctx,
                                                 ggml_tensor* xt, int T, int B,
                                                 ggml_tensor* pe, int pos_len,
                                                 const std::vector<int>& valid_len,
                                                 GraphInputPool& pool,
                                                 int att_left, int att_right) const {
    const int D  = d_model_;
    const int K  = conv_kernel_;
    const float ln_eps = 1e-5f;            // LayerNorm eps (NeMo nn.LayerNorm default)
    const bool local_attn = att_left >= 0;
    assert(local_attn ? (pos_len == att_left + att_right + 1)
                      : (pos_len == 2 * T - 1));

    const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".";
    const ModelLoader& ml = ml_;

    // LayerNorm over the channel dim (ne0 = D), affine. Input ne [D, T, B];
    // ggml_norm normalizes ne0 and the affine [D] broadcasts over T and B.
    auto layer_norm = [&](ggml_tensor* in, const std::string& nm) {
        ggml_tensor* g = clone_weight(ctx, ml, pre + nm + ".weight"); // [D]
        ggml_tensor* b = clone_weight(ctx, ml, pre + nm + ".bias");   // [D]
        ggml_tensor* y = ggml_norm(ctx, in, ln_eps);                  // normalize over ne0
        y = ggml_mul(ctx, y, g);                                      // broadcast [D] over T,B
        y = ggml_add(ctx, y, b);
        return y;
    };
    // nn.Linear: ggml weight ne = [in, out]. in ne [in, T, B] -> [out, T, B].
    auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) {
        ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight");
        ggml_tensor* y = ggml_mul_mat(ctx, W, in);
        if (bias) {
            ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias");
            if (B) y = ggml_add(ctx, y, B);
        }
        return y;
    };
    // ConformerFeedForward: linear1(d->ff) -> SiLU -> linear2(ff->d). in [D, T, B].
    auto feed_forward = [&](ggml_tensor* in, const std::string& ff) {
        ggml_tensor* h = linear(in, ff + ".linear1", /*bias*/true); // [FF, T, B]
        h = ggml_silu(ctx, h);                                      // Swish == SiLU
        h = linear(h, ff + ".linear2", /*bias*/true);               // [D, T, B]
        return h;
    };

    // === Stage A: r = x + 0.5 * FFN1(norm_ff1(x)). ===
    ggml_tensor* h1 = layer_norm(xt, "norm_feed_forward1");
    h1 = feed_forward(h1, "feed_forward1");
    h1 = ggml_scale(ctx, h1, 0.5f);          // fc_factor
    ggml_tensor* r = ggml_add(ctx, xt, h1);  // [D, T, B]

    // === Stage B: r = r + self_attn(norm_self_att(r)). ===
    ggml_tensor* attn_in = layer_norm(r, "norm_self_att");
    RelPosAttention attn(ml_, layer_idx_);
    ggml_tensor* attn_out = local_attn
        ? attn.build_graph_batched_local_chunked(ctx, attn_in, T, B, pe, pos_len, valid_len,
                                                 att_left, att_right, pool)    // [D, T, B]
        : attn.build_graph_batched(ctx, attn_in, T, B, pe, pos_len, valid_len, pool);
    r = ggml_add(ctx, r, attn_out);

    // === Stage C: r = r + conv(norm_conv(r)). ===
    ggml_tensor* c = layer_norm(r, "norm_conv"); // [D, T, B]
    ggml_tensor* conv_out = build_conv_module(ctx, ml, pre, c, D, T, K, B,
                                              valid_len,
                                              conv_norm_type_, conv_causal_, pool);
    r = ggml_add(ctx, r, conv_out);

    // === Stage D: r = r + 0.5 * FFN2(norm_ff2(r)); out = norm_out(r). ===
    ggml_tensor* h2 = layer_norm(r, "norm_feed_forward2");
    h2 = feed_forward(h2, "feed_forward2");
    h2 = ggml_scale(ctx, h2, 0.5f);
    r = ggml_add(ctx, r, h2);
    r = layer_norm(r, "norm_out");
    return r; // [D, T, B] -> per item row-major [T, D]
}

// Scalar (B=1) builder: the verbatim v1 2-D conformer layer graph. The fused
// single-clip encoder and the scalar test entry points (forward /
// forward_with_conv) route through here so B=1 runs the lean 2-D graph and is
// bit-exact with v1. The batched builder above is used by forward_batch (B>1).
ggml_tensor* ConformerLayer::build_graph(ggml_context* ctx, ggml_tensor* xt,
                                         int T, ggml_tensor* pe, int pos_len,
                                         int valid_len,
                                         GraphInputPool& pool,
                                         int att_left, int att_right) const {
    const int D  = d_model_;
    const int K  = conv_kernel_;
    const float ln_eps = 1e-5f;            // LayerNorm eps (NeMo nn.LayerNorm default)
    const bool local_attn = att_left >= 0;
    assert(local_attn ? (pos_len == att_left + att_right + 1)
                      : (pos_len == 2 * T - 1));

    const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".";
    const ModelLoader& ml = ml_;

    // LayerNorm over the channel dim (ne0 = D), affine. Input ne [D, T].
    auto layer_norm = [&](ggml_tensor* in, const std::string& nm) {
        ggml_tensor* g = clone_weight(ctx, ml, pre + nm + ".weight"); // [D]
        ggml_tensor* b = clone_weight(ctx, ml, pre + nm + ".bias");   // [D]
        ggml_tensor* y = ggml_norm(ctx, in, ln_eps);                  // normalize over ne0
        y = ggml_mul(ctx, y, g);                                      // broadcast [D] over T
        y = ggml_add(ctx, y, b);
        return y;
    };
    // nn.Linear: ggml weight ne = [in, out]. in ne [in, T] -> [out, T].
    auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) {
        ggml_tensor* W = clone_weight(ctx, ml, pre + nm + ".weight");
        ggml_tensor* y = ggml_mul_mat(ctx, W, in);
        if (bias) {
            ggml_tensor* B = clone_weight_opt(ctx, ml, pre + nm + ".bias");
            if (B) y = ggml_add(ctx, y, B);
        }
        return y;
    };
    // ConformerFeedForward: linear1(d->ff) -> SiLU -> linear2(ff->d). in [D, T].
    auto feed_forward = [&](ggml_tensor* in, const std::string& ff) {
        ggml_tensor* h = linear(in, ff + ".linear1", /*bias*/true); // [FF, T]
        h = ggml_silu(ctx, h);                                      // Swish == SiLU
        h = linear(h, ff + ".linear2", /*bias*/true);               // [D, T]
        return h;
    };

    // === Stage A: r = x + 0.5 * FFN1(norm_ff1(x)). ===
    ggml_tensor* h1 = layer_norm(xt, "norm_feed_forward1");
    h1 = feed_forward(h1, "feed_forward1");
    h1 = ggml_scale(ctx, h1, 0.5f);          // fc_factor
    ggml_tensor* r = ggml_add(ctx, xt, h1);  // [D, T]

    // === Stage B: r = r + self_attn(norm_self_att(r)). ===
    ggml_tensor* attn_in = layer_norm(r, "norm_self_att");
    RelPosAttention attn(ml_, layer_idx_);
    ggml_tensor* attn_out = local_attn
        ? attn.build_graph_local_chunked(ctx, attn_in, T, pe, pos_len, valid_len,
                                         att_left, att_right, pool)    // [D, T]
        : attn.build_graph(ctx, attn_in, T, pe, pos_len, valid_len, pool); // [D, T]
    r = ggml_add(ctx, r, attn_out);

    // === Stage C: r = r + conv(norm_conv(r)). ===
    ggml_tensor* c = layer_norm(r, "norm_conv"); // [D, T]
    ggml_tensor* conv_out = build_conv_module(ctx, ml, pre, c, D, T, K, valid_len,
                                              conv_norm_type_, conv_causal_, pool);
    r = ggml_add(ctx, r, conv_out);

    // === Stage D: r = r + 0.5 * FFN2(norm_ff2(r)); out = norm_out(r). ===
    ggml_tensor* h2 = layer_norm(r, "norm_feed_forward2");
    h2 = feed_forward(h2, "feed_forward2");
    h2 = ggml_scale(ctx, h2, 0.5f);
    r = ggml_add(ctx, r, h2);
    r = layer_norm(r, "norm_out");
    return r; // [D, T] -> row-major [T, D]
}

void ConformerLayer::forward(const std::vector<float>& x, int T,
                             const std::vector<float>& pos_emb, int pos_len,
                             int valid_len,
                             std::vector<float>& out) const {
    std::vector<float> conv_out_unused;
    forward_with_conv(x, T, pos_emb, pos_len, valid_len, out, conv_out_unused);
}

void ConformerLayer::forward_with_conv(const std::vector<float>& x, int T,
                                       const std::vector<float>& pos_emb, int pos_len,
                                       int valid_len,
                                       std::vector<float>& out,
                                       std::vector<float>& conv_out) const {
    const int D = d_model_;
    assert((int)x.size() == T * D);
    assert((int)pos_emb.size() == pos_len * D);
    assert(pos_len == 2 * T - 1);

    // Unit-test entry point: build the full layer as ONE sub-graph and compute
    // it on the persistent Backend (the fused encoder calls build_graph
    // directly). The conv-module output (for localization) is recomputed via a
    // tiny separate sub-graph below; the test only needs it for a print, not the
    // hot path, so the extra graph is acceptable.
    {
        GraphInputPool pool;
        bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
            [&](ggml_context* ctx) -> ggml_tensor* {
                int64_t xt_ne[2] = {D, T};
                ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2,
                                      xt_ne, x.data(), (size_t)T * D * sizeof(float));
                int64_t pe_ne[2] = {D, pos_len};
                ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2,
                                      pe_ne, pos_emb.data(),
                                      (size_t)pos_len * D * sizeof(float));
                return build_graph(ctx, xt, T, pe, pos_len, valid_len, pool);
            }, out);
        assert(ok && "conformer layer graph failed"); (void)ok;
    }
    // Conv-module output localization (norm_conv -> conv): recompute r up to the
    // conv input and emit the conv-module output. r here = x + FFN1 + attn (same
    // as the fused path's residual into norm_conv).
    {
        GraphInputPool pool;
        bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
            [&](ggml_context* ctx) -> ggml_tensor* {
                int64_t xt_ne[2] = {D, T};
                ggml_tensor* xt = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2,
                                      xt_ne, x.data(), (size_t)T * D * sizeof(float));
                int64_t pe_ne[2] = {D, pos_len};
                ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2,
                                      pe_ne, pos_emb.data(),
                                      (size_t)pos_len * D * sizeof(float));
                const std::string pre = "encoder.layers." +
                                        std::to_string(layer_idx_) + ".";
                auto layer_norm = [&](ggml_tensor* in, const std::string& nm) {
                    ggml_tensor* g = clone_weight(ctx, ml_, pre + nm + ".weight");
                    ggml_tensor* b = clone_weight(ctx, ml_, pre + nm + ".bias");
                    ggml_tensor* y = ggml_norm(ctx, in, 1e-5f);
                    y = ggml_mul(ctx, y, g);
                    y = ggml_add(ctx, y, b);
                    return y;
                };
                auto linear = [&](ggml_tensor* in, const std::string& nm, bool bias) {
                    ggml_tensor* W = clone_weight(ctx, ml_, pre + nm + ".weight");
                    ggml_tensor* y = ggml_mul_mat(ctx, W, in);
                    if (bias) { ggml_tensor* B = clone_weight_opt(ctx, ml_, pre + nm + ".bias");
                                if (B) y = ggml_add(ctx, y, B); }
                    return y;
                };
                auto feed_forward = [&](ggml_tensor* in, const std::string& ff) {
                    ggml_tensor* h = linear(in, ff + ".linear1", true);
                    h = ggml_silu(ctx, h);
                    return linear(h, ff + ".linear2", true);
                };
                ggml_tensor* h1 = feed_forward(layer_norm(xt, "norm_feed_forward1"),
                                               "feed_forward1");
                h1 = ggml_scale(ctx, h1, 0.5f);
                ggml_tensor* r = ggml_add(ctx, xt, h1);
                ggml_tensor* attn_in = layer_norm(r, "norm_self_att");
                RelPosAttention attn(ml_, layer_idx_);
                ggml_tensor* attn_out = attn.build_graph(ctx, attn_in, T, pe,
                                            pos_len, valid_len, pool);
                r = ggml_add(ctx, r, attn_out);
                ggml_tensor* c = layer_norm(r, "norm_conv");
                return build_conv_module(ctx, ml_, pre, c, D, T, conv_kernel_,
                                         valid_len, conv_norm_type_, conv_causal_,
                                         pool);
            }, conv_out);
        assert(ok && "conformer conv localization graph failed"); (void)ok;
    }
}

void ConformerLayer::conv_module_forward(const std::vector<float>& conv_in, int T,
                                         int valid_len,
                                         std::vector<float>& out) const {
    const int D = d_model_;
    const int K = conv_kernel_;
    assert((int)conv_in.size() == T * D);

    const std::string pre = "encoder.layers." + std::to_string(layer_idx_) + ".";
    const ModelLoader& ml = ml_;

    // conv_in is norm_conv(residual), row-major [T, D] -> ggml [D, T]; feed it
    // straight into the conv sub-module (NeMo's ConformerConvolution.forward,
    // which starts AFTER norm_conv).
    GraphInputPool pool;
    bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/4,
        [&](ggml_context* ctx) -> ggml_tensor* {
            int64_t c_ne[2] = {D, T};
            ggml_tensor* c = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, c_ne,
                                 conv_in.data(), (size_t)T * D * sizeof(float));
            return build_conv_module(ctx, ml, pre, c, D, T, K, valid_len,
                                     conv_norm_type_, conv_causal_, pool);
        }, out);
    assert(ok && "conv_module_forward graph failed"); (void)ok;
}

} // namespace pk