transcribe-cpp-sys 0.2.3

// src/conformer/conformer.cpp - shared Conformer encoder helpers.
//
// See conformer.h for the API contract. Each family's encoder.cpp keeps
// only the glue binding these helpers to its weights / hparams / graph
// driver; per-family policy knobs (ConvPolicy, BlockParams) are passed in.
// Conv-dispatch env vars (TRANSCRIBE_CONV_DIRECT_PW/DW and their NO_ variants)
// are read only via resolve_conv_direct here, which detect_direct_pw and every
// family's depthwise detect delegate to — one place owns the parsing.

#include "conformer/conformer.h"

#include "ggml.h"
#include "transcribe-env.h"
#include "transcribe-log.h"

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>

namespace transcribe::conformer {

ggml_tensor * named(ggml_tensor * t, const char * name) {
    if (t != nullptr) {
        ggml_set_name(t, name);
    }
    return t;
}

// LayerNorm with affine: y = gamma * (x - mean) / sqrt(var + eps) + beta.
// ggml_norm normalizes along ne[0] (the d_model axis of [d_model, T, B]);
// gamma/beta [d_model] broadcast over T and B. `gamma` required; `beta`
// may be null for bias-free LN.
ggml_tensor * layer_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * gamma, ggml_tensor * beta) {
    ggml_tensor * y = ggml_norm(ctx, x, kLayerNormEps);
    y               = ggml_mul(ctx, y, gamma);
    if (beta != nullptr) {
        y = ggml_add(ctx, y, beta);
    }
    return y;
}

// FeedForward: y = Linear2(SiLU(Linear1(x))). Weights stored
// ff*_lin1_w ne=[d_model, d_ff], ff*_lin2_w ne=[d_ff, d_model]. Biases
// [d_out] broadcast along T and B; either or both may be null.
ggml_tensor * feed_forward(ggml_context * ctx,
                           ggml_tensor *  x,
                           ggml_tensor *  lin1_w,
                           ggml_tensor *  lin1_b,
                           ggml_tensor *  lin2_w,
                           ggml_tensor *  lin2_b) {
    ggml_tensor * y = ggml_mul_mat(ctx, lin1_w, x);
    if (lin1_b != nullptr) {
        y = ggml_add(ctx, y, lin1_b);
    }
    y = ggml_silu(ctx, y);
    y = ggml_mul_mat(ctx, lin2_w, y);
    if (lin2_b != nullptr) {
        y = ggml_add(ctx, y, lin2_b);
    }
    return y;
}

// Macaron half-residual feed-forward: x + 0.5 * FF(LN(x)). The 0.5 scale
// is the macaron weighting (FF1/FF2); MHSA and conv use a full residual.
ggml_tensor * macaron_ff_residual(ggml_context * ctx,
                                  ggml_tensor *  x,
                                  ggml_tensor *  norm_w,
                                  ggml_tensor *  norm_b,
                                  ggml_tensor *  lin1_w,
                                  ggml_tensor *  lin1_b,
                                  ggml_tensor *  lin2_w,
                                  ggml_tensor *  lin2_b) {
    ggml_tensor * y = layer_norm(ctx, x, norm_w, norm_b);
    y               = feed_forward(ctx, y, lin1_w, lin1_b, lin2_w, lin2_b);
    y               = ggml_scale(ctx, y, 0.5f);
    return ggml_add(ctx, x, y);
}

// Shaw / Transformer-XL relative-position skew. Rotates each row of the
// [pos_len, T_q, H, 1] score matrix so column k holds relative offset k.
ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x) {
    const int64_t pos_len = x->ne[0];
    const int64_t T_q     = x->ne[1];
    const int64_t H       = x->ne[2];
    const int64_t B       = x->ne[3];

    // Step 1: prepend a [1, T_q, H, B] column of zeros along axis 0.
    // ggml_fill takes a shape template (need not be backed by real
    // data — only the metadata is used) and emits a new tensor of
    // the same shape filled with the given constant. ggml_concat
    // along dim=0 produces [pos_len+1, T_q, H, B] with the first
    // axis-0 element being our zero column.
    ggml_tensor * zero_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, T_q, H, B);
    ggml_tensor * zeros         = ggml_fill(ctx, zero_template, 0.0f);
    ggml_tensor * y             = ggml_concat(ctx, zeros, x, /*dim=*/0);

    // Step 2: reshape (swap inner two axes; total elements unchanged).
    y = ggml_reshape_4d(ctx, y, T_q, pos_len + 1, H, B);

    // Step 3: drop the first row of axis 1.
    y = ggml_view_4d(ctx, y, T_q, pos_len, H, B, y->nb[1], y->nb[2], y->nb[3],
                     /*offset=*/y->nb[1]);

    // Step 4: materialize as contiguous (the view above is sparse:
    // its nb[1] still reflects the parent's stride which assumes
    // pos_len+1 elements per row, so reshape can't read it directly).
    y = ggml_cont(ctx, y);

    // Step 5: reshape back to [pos_len, T_q, H, B].
    y = ggml_reshape_4d(ctx, y, pos_len, T_q, H, B);
    return y;
}

// f32-friendly Conv1D (mirrors ggml_conv_1d but passes the kernel's real
// type to im2col instead of forcing f16). Vendored ggml's ggml_conv_1d
// hardcodes GGML_TYPE_F16 for the im2col output and asserts on an f32 kernel.
ggml_tensor * conv_1d_f32(ggml_context * ctx,
                          ggml_tensor *  kernel,
                          ggml_tensor *  data,
                          int            stride,
                          int            padding,
                          int            dilation) {
    // im2col: kernel shape [K, IC, OC] in ggml ne; data shape
    // [W, IC, N] in ggml ne. Output ne[0] = IC*K, ne[1] = OW (output
    // width), ne[2] = N (batch).
    ggml_tensor * im2col = ggml_im2col(ctx, kernel, data, stride, /*s1=*/0, padding, /*p1=*/0, dilation, /*d1=*/0,
                                       /*is_2D=*/false,
                                       /*dst_type=*/kernel->type);
    // im2col ne = [IC*K, OW, N].

    const int64_t N         = im2col->ne[2];
    ggml_tensor * kernel_2d = ggml_reshape_2d(ctx, kernel, kernel->ne[0] * kernel->ne[1],
                                              kernel->ne[2]);  // [IC*K, OC]

    // F32 accumulation when the kernel is F16 (see mul_mat_f32acc in
    // causal_lm.cpp for the CUDA COMPUTE_16F saturation rationale).
    const bool kernel_needs_f32_acc = (kernel->type == GGML_TYPE_F16);

    if (N == 1) {
        // Single-shot path: flatten OW*N (== OW) and reshape the [OW, OC]
        // result back to [OW, OC, 1].
        ggml_tensor * result = ggml_mul_mat(ctx, ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[1]),
                                            kernel_2d);  // [OW, OC]
        if (kernel_needs_f32_acc) {
            ggml_mul_mat_set_prec(result, GGML_PREC_F32);
        }
        result = ggml_reshape_3d(ctx, result, im2col->ne[1], kernel->ne[2], 1);
        return result;
    }

    // Batched path. Flattening OW*N interleaves batch with width
    // (ow_n = b*OW + ow), so a [OW*N, OC] -> [OW, OC, N] reshape would
    // mislabel memory. Instead mul_mat the 3-D im2col directly: the kernel
    // (ne[2]=1) broadcasts across the batch, giving [OC, OW, N], then
    // permute to the [OW, OC, N] = [time, channels, batch] convention.
    ggml_tensor * result = ggml_mul_mat(ctx, kernel_2d, im2col);  // [OC, OW, N]
    if (kernel_needs_f32_acc) {
        ggml_mul_mat_set_prec(result, GGML_PREC_F32);
    }
    result = ggml_cont(ctx, ggml_permute(ctx, result, 1, 0, 2, 3));
    return result;  // [OW, OC, N]
}

// f32-friendly 2D depthwise conv (mirrors ggml_conv_2d_dw but passes the
// kernel's real type to im2col). CANONICAL Metal conv-quirk note: vendored
// ggml's ggml_conv_2d_dw hardcodes F16 im2col, which routes the matmul to
// a `kernel_mul_mv_f32_f16_short` Metal kernel that isn't compiled in;
// forcing f32 lands on the available `kernel_mul_mv_f32_f32_short`. The
// alternative ggml_conv_2d_dw_direct can't help: its CONV_2D_DW op is
// unsupported on Metal, so the im2col + mul_mat path here is what works.
ggml_tensor * conv_2d_dw_f32(ggml_context * ctx,
                             ggml_tensor *  kernel,  // [KW, KH, 1, C]
                             ggml_tensor *  data,    // [W, H, C, N]
                             int            s0,
                             int            s1,
                             int            p0,
                             int            p1,
                             int            d0,
                             int            d1) {
    // Same op sequence as ggml_conv_2d_dw, only dst_type changes.
    ggml_tensor * new_a = ggml_reshape_4d(ctx, kernel, kernel->ne[0], kernel->ne[1], 1, kernel->ne[2] * kernel->ne[3]);
    ggml_tensor * data_4d = ggml_reshape_4d(ctx, data, data->ne[0], data->ne[1], 1, data->ne[2] * data->ne[3]);
    ggml_tensor * im2col  = ggml_im2col(ctx, new_a, data_4d, s0, s1, p0, p1, d0, d1,
                                        /*is_2D=*/true,
                                        /*dst_type=*/kernel->type);
    ggml_tensor * new_b =
        ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], data->ne[2], data->ne[3]);
    new_a                = ggml_reshape_4d(ctx, new_a, new_a->ne[0] * new_a->ne[1], new_a->ne[2], new_a->ne[3], 1);
    ggml_tensor * result = ggml_mul_mat(ctx, new_a, new_b);
    result               = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], data->ne[2], data->ne[3]);
    return result;
}

// ggml_conv_2d_dw_direct requires an F32 kernel on CUDA (a hard GGML_ASSERT
// in conv2d-dw.cu) and silently misbehaves for F16 kernels on other backends
// — see conv_1d_dw_f32's batch path below and the canary_qwen note. A
// BF16-reference conformer (e.g. cohere) ships its depthwise kernels at F16,
// so promote to F32 in-graph before any direct depthwise op. The depthwise
// kernel is tiny, so the cast is negligible; F32 kernels pass through
// untouched. Only the direct_dw path (CUDA/Vulkan) routes here — the im2col
// fallback used on Metal/CPU takes the kernel's real type and is unaffected.
static ggml_tensor * dw_kernel_for_direct(ggml_context * ctx, ggml_tensor * kernel) {
    if (kernel == nullptr || kernel->type == GGML_TYPE_F32) {
        return kernel;
    }
    return ggml_cast(ctx, kernel, GGML_TYPE_F32);
}

// f32-friendly depthwise Conv1D (mirrors ggml_conv_1d_dw but
// passes the kernel's real type to im2col). Same fix as above.
ggml_tensor * conv_1d_dw_f32(ggml_context * ctx,
                             ggml_tensor *  kernel,
                             ggml_tensor *  data,
                             int            stride,
                             int            padding,
                             int            dilation) {
    // Offline batch: data ne=[W, C, B] with B>1. The im2col path below
    // collapses the batch axis (its final reshape hardcodes ne[2]=1), so
    // route B>1 through the direct depthwise-2D op (W, H=1, C, N=B), which
    // threads the utterance batch at ne[3].
    const int64_t B = data->ne[2];
    if (B > 1) {
        const int64_t k   = kernel->ne[0];
        const int64_t C   = data->ne[1];
        // ggml_conv_2d_dw_direct misbehaves for non-f32 kernels; the
        // depthwise kernel is tiny, so cast to f32.
        ggml_tensor * knl = kernel;
        if (knl->type != GGML_TYPE_F32) {
            knl = ggml_cast(ctx, knl, GGML_TYPE_F32);
        }
        knl              = ggml_reshape_4d(ctx, knl, k, 1, 1, C);
        ggml_tensor * d4 = ggml_reshape_4d(ctx, data, data->ne[0], 1, C, B);
        ggml_tensor * o  = ggml_conv_2d_dw_direct(ctx, knl, d4,
                                                  /*s0=*/stride, /*s1=*/1,
                                                  /*p0=*/padding, /*p1=*/0,
                                                  /*d0=*/dilation, /*d1=*/1);
        return ggml_reshape_3d(ctx, o, o->ne[0], C, B);  // [W_out, C, B]
    }

    // The depthwise wrapper reshapes the data into a 2D image with
    // H=1 so the same im2col path can be reused. ggml_conv_1d_dw
    // does this; we replicate it here with the dst_type override.
    ggml_tensor * data_4d = ggml_reshape_4d(ctx, data, data->ne[0], 1, data->ne[1], data->ne[2]);
    ggml_tensor * im2col  = ggml_im2col(ctx, kernel, data_4d, stride, /*s1=*/0, padding, /*p1=*/0, dilation, /*d1=*/0,
                                        /*is_2D=*/false,
                                        /*dst_type=*/kernel->type);
    ggml_tensor * result  = ggml_mul_mat(ctx, im2col, kernel);
    result                = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1);
    return result;
}

// Fused BatchNorm: y = x * scale + bias, where scale and bias are
// precomputed at load time from the raw BN parameters. Replaces the
// 4-op batch_norm (sub, div/sqrt, mul, add) with 2 ops (mul, add).
// The 1-D tensors [d_model] are reshaped to [1, d_model, 1, 1] to
// broadcast across the time axis of [T, d_model, 1, 1].
ggml_tensor * fused_batch_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale_1d, ggml_tensor * bias_1d) {
    const int64_t C        = scale_1d->ne[0];
    ggml_tensor * scale_4d = ggml_reshape_4d(ctx, scale_1d, 1, C, 1, 1);
    ggml_tensor * bias_4d  = ggml_reshape_4d(ctx, bias_1d, 1, C, 1, 1);
    ggml_tensor * y        = ggml_mul(ctx, x, scale_4d);
    y                      = ggml_add(ctx, y, bias_4d);
    return y;
}

// Add a 1-D bias [C] to a 4-D conv output [W, H, C, N] by reshaping
// the bias to [1, 1, C, 1] so ggml_add broadcasts the W, H, and N
// axes for free. The reshape is a metadata-only view because the
// 1-D tensor is contiguous; no data copy.
ggml_tensor * add_conv_bias(ggml_context * ctx, ggml_tensor * conv_out, ggml_tensor * bias_1d) {
    if (bias_1d == nullptr) {
        return conv_out;
    }
    const int64_t channels = bias_1d->ne[0];
    ggml_tensor * bias_4d  = ggml_reshape_4d(ctx, bias_1d, 1, 1, channels, 1);
    return ggml_add(ctx, conv_out, bias_4d);
}

bool resolve_conv_direct(const char * direct_env, const char * no_direct_env, bool backend_default) {
    if (transcribe::env::flag(direct_env)) {
        return true;  // user override
    }
    if (transcribe::env::flag(no_direct_env)) {
        return false;  // user override
    }
    return backend_default;
}

bool detect_direct_pw(const char * backend) {
    // Vulkan defaults to im2col: on AMD Renoir, im2col + mul_mat measured
    // ~200 ms/encode faster than direct for f32 weights. Metal/CPU prefer
    // direct.
    const bool backend_default = !(backend != nullptr && std::strstr(backend, "Vulkan") != nullptr);
    return resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_PW", "TRANSCRIBE_CONV_NO_DIRECT_PW", backend_default);
}

// Conv module (pointwise -> GLU -> depthwise -> BN/LN -> SiLU -> pointwise).
// Operates on the post-LayerNorm activation (LN applied by the caller).
// Pointwise convs (k=1) are direct mul_mats in [d_model, T] layout by
// default; TRANSCRIBE_CONV_NO_DIRECT_PW=1 forces the im2col path. BatchNorm
// uses fused scale + bias; LayerNorm normalizes per-channel on the fly.
ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & b, const BlockParams & params) {
    const int          conv_kernel = params.conv_kernel;
    const ConvPolicy & policy      = params.policy;

    // Centred (k-1)/2 by default; (conv_context_left, conv_context_right)
    // exactly when both are >= 0 (caller asked for causal or otherwise
    // asymmetric padding).
    int pad_left  = (conv_kernel - 1) / 2;
    int pad_right = (conv_kernel - 1) / 2;
    if (params.conv_context_left >= 0 && params.conv_context_right >= 0) {
        pad_left  = params.conv_context_left;
        pad_right = params.conv_context_right;
    }

    const int64_t d_model = x->ne[0];
    // Utterance batch lives at ne[2] of the [d_model, T, B] activation.
    // B == 1 for single-shot; the offline batched encoder passes B > 1.
    const int64_t B       = x->ne[2];

    if (policy.direct_pw) {
        // Pointwise conv 1 as direct mul_mat in [d_model, T, B] layout.
        // Kernel ne=[1, d_model, 2*d_model] → reshape to [d_model, 2*d_model].
        // Force F32 accumulation for F16 weights (see mul_mat_f32acc in
        // causal_lm.cpp for the CUDA COMPUTE_16F saturation rationale).
        {
            ggml_tensor * pw1 = ggml_reshape_2d(ctx, b.conv_pw1_w, d_model, 2 * d_model);
            x                 = ggml_mul_mat(ctx, pw1, x);  // [2*d_model, T, B]
            if (b.conv_pw1_w->type == GGML_TYPE_F16) {
                ggml_mul_mat_set_prec(x, GGML_PREC_F32);
            }
            if (b.conv_pw1_b != nullptr) {
                x = ggml_add(ctx, x, b.conv_pw1_b);
            }
        }

        // GLU: split ne[0] in half, gate * sigmoid(value). The views carry
        // the batch axis (ne[2]) so the split is per-utterance.
        {
            const int64_t T     = x->ne[1];
            const int64_t half  = x->ne[0] / 2;
            ggml_tensor * gate  = ggml_view_3d(ctx, x, half, T, B, x->nb[1], x->nb[2],
                                               /*offset=*/0);
            ggml_tensor * value = ggml_view_3d(ctx, x, half, T, B, x->nb[1], x->nb[2], half * ggml_element_size(x));
            x                   = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));
        }
        // x ne = [d_model, T, B]

        // Transpose for depthwise conv: [d_model, T, B] -> [T, d_model, B].
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
    } else {
        // Fallback: im2col path in [T, d_model] layout.
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

        x = conv_1d_f32(ctx, b.conv_pw1_w, x, /*s=*/1, /*p=*/0, /*d=*/1);

        // Add pointwise1 bias in [T, 2*d_model] layout.
        if (b.conv_pw1_b != nullptr) {
            // conv_1d_f32 returns ne=[T, 2*d_model, 1]. Reshape the
            // 1-D bias to [1, 2*d_model] so ggml_add broadcasts.
            x                    = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1]);
            ggml_tensor * bias_r = ggml_reshape_2d(ctx, b.conv_pw1_b, 1, 2 * d_model);
            x                    = ggml_add(ctx, x, bias_r);
        }

        const int64_t half  = x->ne[1] / 2;
        ggml_tensor * gate  = ggml_view_4d(ctx, x, x->ne[0], half, 1, 1, x->nb[1], x->nb[2], x->nb[3], 0);
        ggml_tensor * value = ggml_view_4d(ctx, x, x->ne[0], half, 1, 1, x->nb[1], x->nb[2], x->nb[3], x->nb[1] * half);
        x                   = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value));

        x = ggml_cont(ctx, x);
    }

    // NeMo zeros padded-overhang frames after pointwise_conv1 + GLU and
    // before the depthwise convolution. Without this, buffered streaming
    // overhang frames can leak backward through the conv kernel.
    if (params.conv_pad_mask != nullptr) {
        x = ggml_mul(ctx, x, params.conv_pad_mask);
    }

    // Depthwise conv: kernel size from hparams, groups=d_model. The
    // ggml depthwise ops accept a single symmetric padding value on
    // each axis; for asymmetric padding (causal: [k-1, 0]) we prepend
    // / append zero-frames along the time axis explicitly and call the
    // op with p0=0. At this point x has ne = [T, d_model, 1, 1] (the
    // permute above moved T into ne[0]), so we concat along dim 0.
    //
    // Streaming carry: when params.streaming_time_in is set, the left
    // pad is the previous chunk's last (k-1) post-pw1+GLU frames
    // (NeMo's cache_last_time) instead of zeros. The new cache slot
    // for the next chunk is the LAST pad_left frames of the
    // concatenated (prev_cache, current_x) sequence — NeMo's
    //   cache_next = concat(prev_cache, x)[-pad_left:]
    // which works regardless of whether T_now >= pad_left. The
    // earlier "tail of x" formulation broke at small chunk sizes
    // (R<13 in nemotron-streaming produces T_q_new < pad_left=8).
    const bool symmetric_pad = (pad_left == pad_right);
    if (!symmetric_pad) {
        const bool streaming = (params.streaming_time_in != nullptr);

        if (streaming) {
            if (pad_left > 0) {
                // streaming_time_in has ne = [pad_left, d_model, 1, 1]
                // (matches the zero-pad shape on this axis).
                x = ggml_concat(ctx, params.streaming_time_in, x, /*dim=*/0);
            }

            // After the concat, x ne = [pad_left + T_now, d_model, 1, 1].
            // Take the LAST pad_left frames as the next-chunk cache slot.
            // pad_left always <= x->ne[0] here because the concat above
            // contributes exactly pad_left frames of prev_cache.
            if (params.streaming_time_out != nullptr && params.streaming_graph != nullptr) {
                const int64_t T_padded = x->ne[0];
                ggml_tensor * tail = ggml_view_4d(ctx, x, pad_left, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2],
                                                  x->nb[3], (T_padded - pad_left) * x->nb[0]);
                ggml_tensor * cpy  = ggml_cpy(ctx, tail, params.streaming_time_out);
                // Cache writes are SIDE outputs of the encoder graph:
                // not reachable from `eb.out`, so the scheduler won't
                // schedule them unless we expand them explicitly.
                ggml_build_forward_expand(params.streaming_graph, cpy);
            }
        } else if (pad_left > 0) {
            ggml_tensor * pad_l = ggml_new_tensor_4d(ctx, x->type, pad_left, x->ne[1], x->ne[2], x->ne[3]);
            pad_l               = ggml_fill(ctx, pad_l, 0.0f);
            x                   = ggml_concat(ctx, pad_l, x, /*dim=*/0);
        }
        if (pad_right > 0) {
            ggml_tensor * pad_r = ggml_new_tensor_4d(ctx, x->type, pad_right, x->ne[1], x->ne[2], x->ne[3]);
            pad_r               = ggml_fill(ctx, pad_r, 0.0f);
            x                   = ggml_concat(ctx, x, pad_r, /*dim=*/0);
        }
    }
    const int padding_op = symmetric_pad ? pad_left : 0;
    if (policy.direct_dw_in_block) {
        // Fused single-op depthwise conv (no im2col). Input is [T, d_model, B]
        // from the transpose above. Reshape to 4D [T, 1, d_model, B] for
        // ggml_conv_2d_dw_direct which expects [W, H, C, N] (N == B, the
        // utterance batch). Kernel [k, 1, d_model] → [k, 1, 1, d_model].
        ggml_tensor * knl  = ggml_reshape_4d(ctx, dw_kernel_for_direct(ctx, b.conv_dw_w), conv_kernel, 1, 1, d_model);
        ggml_tensor * data = ggml_reshape_4d(ctx, x, x->ne[0], 1, x->ne[1], B);
        x                  = ggml_conv_2d_dw_direct(ctx, knl, data,
                                                    /*s0=*/1, /*s1=*/1,
                                                    /*p0=*/padding_op, /*p1=*/0,
                                                    /*d0=*/1, /*d1=*/1);
        // Output: [T_out, 1, d_model, B] → [T_out, d_model, B].
        x                  = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[2], x->ne[3]);
    } else {
        // im2col depthwise (cohere on Metal/CPU); single-utterance only.
        x = conv_1d_dw_f32(ctx, b.conv_dw_w, x,
                           /*s=*/1, /*p=*/padding_op, /*d=*/1);
    }

    // Add depthwise bias in [T, d_model] layout (nullable).
    if (b.conv_dw_b != nullptr) {
        ggml_tensor * bias_r = ggml_reshape_2d(ctx, b.conv_dw_b, 1, d_model);
        x                    = ggml_add(ctx, x, bias_r);
    }

    // Post-depthwise normalisation: BN (fused mul + add) or LN (per-channel
    // mean/std + affine). x is [T, d_model, 1, 1] here; `layer_norm` reads
    // ne[0] as the feature axis, so for LN we permute to [d_model, T], apply,
    // and permute back.
    if (params.conv_norm_type == BlockParams::ConvNormType::LayerNorm) {
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
        x = layer_norm(ctx, x, b.conv_ln_w, b.conv_ln_b);
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
    } else {
        x = fused_batch_norm(ctx, x, b.conv_bn_fused_scale, b.conv_bn_fused_bias);
    }

    x = ggml_silu(ctx, x);

    if (policy.direct_pw) {
        // Transpose back: [T, d_model] -> [d_model, T].
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

        // Pointwise conv 2 as direct mul_mat in [d_model, T] layout. See
        // the pw1 comment above for the F16 / CUDA COMPUTE_16F rationale.
        ggml_tensor * pw2 = ggml_reshape_2d(ctx, b.conv_pw2_w, d_model, d_model);
        x                 = ggml_mul_mat(ctx, pw2, x);
        if (b.conv_pw2_w->type == GGML_TYPE_F16) {
            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
        }
        if (b.conv_pw2_b != nullptr) {
            x = ggml_add(ctx, x, b.conv_pw2_b);
        }
    } else {
        x = conv_1d_f32(ctx, b.conv_pw2_w, x, /*s=*/1, /*p=*/0, /*d=*/1);
        if (b.conv_pw2_b != nullptr) {
            x                    = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1]);
            ggml_tensor * bias_r = ggml_reshape_2d(ctx, b.conv_pw2_b, 1, d_model);
            x                    = ggml_add(ctx, x, bias_r);
        }
        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
    }

    return x;
}

// Relative-position multi-head self-attention.
//
// use_flash = true  -> ggml_flash_attn_ext (fused GPU kernel, needs
//                      a dk template in the Metal backend)
// use_flash = false -> manual mul_mat + soft_max_ext + mul_mat (works
//                      on any backend; use this when the head dim has
//                      no flash kernel, e.g. Cohere's encoder dk=160)
//
// ggml_flash_attn_ext expected layouts:
//   Q:    [head_dim, T_q,    n_head, 1]
//   K:    [head_dim, T_k,    n_head, 1]
//   V:    [head_dim, T_k,    n_head, 1]  (NOT transposed)
//   mask: [T_k,      T_q,    n_head, 1]  F16, additive
//   out:  [head_dim, n_head, T_q,    1]  (already permuted)
//
// Score formula:
//   flash_attn computes: softmax(q_u @ k^T * scale + mask) @ v
//   We set mask = rel_shift(q_v @ p^T)[:T_k] * scale
//   Which gives: softmax((q_u @ k^T + rel_pos_bias) * scale) @ v
ggml_tensor * rel_pos_mhsa(ggml_context *      ctx,
                           ggml_tensor *       x,
                           ggml_tensor *       pos_emb,
                           const BlockView &   b,
                           const BlockParams & params,
                           ggml_tensor *       x_q,
                           ggml_tensor *       k_full,
                           ggml_tensor *       v_full) {
    const int       d_model           = params.d_model;
    const int       n_head            = params.n_head;
    const ggml_type kv_type           = params.kv_type;
    const bool      use_flash         = params.use_flash;
    const int       att_context_left  = params.att_context_left;
    const int       att_context_right = params.att_context_right;
    const int       head_dim          = d_model / n_head;
    const float     scale             = 1.0f / std::sqrt(static_cast<float>(head_dim));
    // Query/key-value split (see the x_q / k_full contracts in conformer.h).
    // kv_cached: K/V arrive pre-projected, x is the query-only activation.
    // q_sliced: queries from x_q, K/V projected from x. rect covers both —
    // the score geometry is rectangular [T_kv, T_q].
    const bool      kv_cached         = (k_full != nullptr && v_full != nullptr);
    const bool      q_sliced          = !kv_cached && (x_q != nullptr && x_q != x);
    const bool      rect              = kv_cached || q_sliced;
    const int64_t   T_kv              = kv_cached ? k_full->ne[1] : x->ne[1];
    const int64_t   T_q               = q_sliced ? x_q->ne[1] : x->ne[1];
    // With a precomputed pos projection, pos_emb may be null (the
    // caller's graph carries no pos_emb input at all).
    ggml_tensor *   pos_proj          = params.streaming_pos_proj_in;
    const int64_t   pos_len           = pos_proj != nullptr ? pos_proj->ne[1] : pos_emb->ne[1];
    if (rect && pos_len != T_q + T_kv - 1) {
        std::fprintf(stderr,
                     "conformer rel_pos_mhsa: rectangular pos_emb length "
                     "%lld != T_q + T_kv - 1 = %lld\n",
                     (long long) pos_len, (long long) (T_q + T_kv - 1));
        return nullptr;
    }
    // Utterance batch at ne[2] (moves to ne[3] after the head split). Flash
    // works batched (ggml_flash_attn_ext accepts the per-utterance rel-pos
    // mask at ne[3] == B). Rectangular geometry forces manual: flash requires
    // the mask's ne[1] padded to GGML_KQ_MASK_PAD, which the [T_kv, T_q]
    // streaming mask does not satisfy. (TRANSCRIBE_NO_FLASH=1 also forces
    // manual, for the bit-exact CPU tensor gate.)
    const int64_t B     = x->ne[2];
    const bool    flash = use_flash && !rect;

    // Local-attention bookkeeping. With both window sides non-negative
    // in the Regular style, pos_emb arrives at the smaller
    // [left+right+1, d] length and matrix_bd needs zero/-inf padding to
    // keep the rel_shift trick intact. ChunkedLimited keeps pos_emb at
    // its full 2T-1 length and uses the external chunked mask instead.
    const bool is_chunked = (params.att_context_style == BlockParams::AttContextStyle::ChunkedLimited);
    const bool is_local   = (!is_chunked) && (att_context_left >= 0 && att_context_right >= 0);
    const int  W_left     = is_local ? att_context_left : 0;
    const int  W_right    = is_local ? att_context_right : 0;

    // Q, K, V, P projections. Q/K/V may have bias (Cohere) or not
    // (Parakeet); the pos projection (attn_pos_w) never has a bias.
    ggml_tensor * q = ggml_mul_mat(ctx, b.attn_q_w, q_sliced ? x_q : x);
    if (b.attn_q_b != nullptr) {
        q = ggml_add(ctx, q, b.attn_q_b);
    }
    // KV-cache mode: keys/values arrive pre-projected (bias included).
    ggml_tensor * k = k_full;
    ggml_tensor * v = v_full;
    if (!kv_cached) {
        k = ggml_mul_mat(ctx, b.attn_k_w, x);
        if (b.attn_k_b != nullptr) {
            k = ggml_add(ctx, k, b.attn_k_b);
        }
        v = ggml_mul_mat(ctx, b.attn_v_w, x);
        if (b.attn_v_b != nullptr) {
            v = ggml_add(ctx, v, b.attn_v_b);
        }
    }
    ggml_tensor * p = pos_proj == nullptr ? ggml_mul_mat(ctx, b.attn_pos_w, pos_emb) : nullptr;

    // Split heads: [head_dim, n_head, T, B] (batch at ne[3]). pos_bias_u/v
    // broadcast onto this BEFORE the permute that moves T past n_head.
    q                 = ggml_reshape_4d(ctx, q, head_dim, n_head, T_q, B);
    ggml_tensor * q_u = ggml_add(ctx, q, b.attn_pos_u);
    ggml_tensor * q_v = ggml_add(ctx, q, b.attn_pos_v);

    // All of Q, K, V need [head_dim, T, n_head, B] for the attention. The
    // permute is a zero-cost view (no data copy). flash_attn_ext accesses
    // data via strides, so contiguous materialization (cont) is not
    // required for q_u/k/v. q_v and p still need cont because they feed
    // into mul_mat for the position score computation.
    q_u = ggml_permute(ctx, q_u, 0, 2, 1, 3);
    q_v = ggml_cont(ctx, ggml_permute(ctx, q_v, 0, 2, 1, 3));

    k = ggml_reshape_4d(ctx, k, head_dim, n_head, T_kv, B);
    k = ggml_permute(ctx, k, 0, 2, 1, 3);

    v = ggml_reshape_4d(ctx, v, head_dim, n_head, T_kv, B);
    v = ggml_permute(ctx, v, 0, 2, 1, 3);

    // Position scores are batch-independent (pos_emb has no batch axis), so
    // p keeps ne[3] == 1 and broadcasts across the batch in the mul_mat.
    // The precomputed pos_proj is exactly this tensor, memoized.
    if (p != nullptr) {
        p = ggml_reshape_4d(ctx, p, head_dim, n_head, pos_len, 1);
        p = ggml_cont(ctx, ggml_permute(ctx, p, 0, 2, 1, 3));
    } else {
        p = pos_proj;
    }

    // Position mask / bias: matrix_bd = rel_shift(q_v @ p^T), truncated.
    ggml_tensor * matrix_bd = ggml_mul_mat(ctx, p, q_v);

    // Local-attention pad/slice. The standard rel_shift trick assumes
    // matrix_bd has shape [2T_q-1, T_q]: row r corresponds to relative
    // offset (T_q-1-r). For local attention pos_emb is shorter
    // ([W_left+W_right+1]) where row r corresponds to offset (W_left-r).
    // Bring matrix_bd back to the [2T_q-1, T_q] shape by:
    //   - prepending (T_q-1-W_left) rows of -INF (or slicing them off
    //     when the audio is so short the window already covers it),
    //   - appending  (T_q-1-W_right) rows of -INF (or slicing).
    // After this, rel_shift + the existing T_q×T_q view land each
    // out-of-window position at -INF, which softmax zeroes out. With
    // both window sides == -1 (full attention) this block is skipped.
    if (is_local) {
        const int top_pad = static_cast<int>(T_q) - 1 - W_left;
        if (top_pad > 0) {
            ggml_tensor * top_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, top_pad, T_q, n_head, B);
            ggml_tensor * top          = ggml_fill(ctx, top_template, -INFINITY);
            matrix_bd                  = ggml_concat(ctx, top, matrix_bd, /*dim=*/0);
        } else if (top_pad < 0) {
            const int kept = static_cast<int>(matrix_bd->ne[0]) + top_pad;
            matrix_bd      = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2],
                                          matrix_bd->nb[3], (-top_pad) * matrix_bd->nb[0]);
            matrix_bd      = ggml_cont(ctx, matrix_bd);
        }
        const int bot_pad = static_cast<int>(T_q) - 1 - W_right;
        if (bot_pad > 0) {
            ggml_tensor * bot_template = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, bot_pad, T_q, n_head, B);
            ggml_tensor * bot          = ggml_fill(ctx, bot_template, -INFINITY);
            matrix_bd                  = ggml_concat(ctx, matrix_bd, bot, /*dim=*/0);
        } else if (bot_pad < 0) {
            const int kept = static_cast<int>(matrix_bd->ne[0]) + bot_pad;
            matrix_bd      = ggml_view_4d(ctx, matrix_bd, kept, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2],
                                          matrix_bd->nb[3], /*offset=*/0);
            matrix_bd      = ggml_cont(ctx, matrix_bd);
        }
    }

    // rel_shift generalizes to the rectangular case: with input
    // [T_q + T_kv - 1, T_q] it yields out[k, q] = in[k - q + T_q - 1, q],
    // i.e. row k holds the score of key k against query q for relative
    // offset (T_kv - 1) - (k - q + T_q - 1) = (T_kv - T_q) + q - k —
    // exactly the query-at-absolute-position (T_kv - T_q + q) semantics
    // the streaming x_q path needs. The square offline case is the
    // T_q == T_kv specialization. The zero column injected by the trick
    // only lands at k >= T_kv, which the view below slices off.
    matrix_bd = rel_shift(ctx, matrix_bd);
    matrix_bd = ggml_view_4d(ctx, matrix_bd, T_kv, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], matrix_bd->nb[3],
                             /*offset=*/0);
    // The view is non-contiguous (nb[1] stays at parent's
    // pos_len*es), but it IS contiguous-rows. The flash path calls
    // ggml_scale which wants full contiguity, so cont there. The
    // manual path only feeds matrix_bd into ggml_add(kq, matrix_bd),
    // which handles contiguous-rows inputs on both CPU and Metal.
    if (flash) {
        matrix_bd = ggml_cont(ctx, matrix_bd);
    }

    // ChunkedLimited mask. Caller provided a [T_q, T_q, 1, 1] F32
    // tensor with 0 on allowed (q, k) pairs and -INF outside the
    // [q_chunk - left_chunks, q_chunk] band. Broadcasts across n_head.
    // -INF and 0 are scale-invariant so this can be added before the
    // pre-scale that the flash path applies below.
    if (is_chunked && params.attn_chunked_mask != nullptr) {
        matrix_bd = ggml_add(ctx, matrix_bd, params.attn_chunked_mask);
    }

    // Variable-length batch key-padding mask. [T_k, 1, 1, B] additive
    // (-INF on padded keys) broadcasts over queries and heads. Added here
    // so it applies on both the flash and manual paths (matrix_bd is the
    // flash mask and the manual additive bias alike). -INF / 0 are
    // scale-invariant, so adding before the flash pre-scale is fine.
    if (params.attn_pad_mask != nullptr) {
        matrix_bd = ggml_add(ctx, matrix_bd, params.attn_pad_mask);
    }

    ggml_tensor * o;

    if (flash) {
        matrix_bd = ggml_scale(ctx, matrix_bd, scale);
        matrix_bd = ggml_cast(ctx, matrix_bd, GGML_TYPE_F16);

        // Optionally cast K/V activations to a narrower type to
        // reduce bandwidth in the attention kernel. GGML_TYPE_COUNT
        // means "auto" (f16 for quantized weights, f32 for f32
        // weights). The flash_attn accumulator is f32 internally
        // regardless of K/V type.
        {
            ggml_type effective_kv = kv_type;
            if (effective_kv == GGML_TYPE_COUNT) {
                effective_kv = (b.attn_k_w->type != GGML_TYPE_F32) ? GGML_TYPE_F16 : GGML_TYPE_F32;
            }
            if (effective_kv != GGML_TYPE_F32) {
                k = ggml_cast(ctx, k, effective_kv);
                v = ggml_cast(ctx, v, effective_kv);
            }
        }

        o = ggml_flash_attn_ext(ctx, q_u, k, v, matrix_bd, scale, /*max_bias=*/0.0f,
                                /*logit_softcap=*/0.0f);
        // Output is [head_dim, n_head, T_q, 1] — already permuted.
    } else {
        // Manual attention path: explicit matmuls + softmax.
        //
        //   attn_scores  = q_u @ k^T                    [T_q, T_q, n_head, 1]
        //   attn_scores  = (attn_scores + matrix_bd) * scale
        //   attn_weights = softmax(attn_scores)
        //   o            = attn_weights @ v
        //
        // K is [head_dim, T_q, n_head, 1] from permute(0,2,1,3).
        // mul_mat(K, Q_u) computes Q_u^T @ K for each head, giving
        // [T_q, T_q, n_head, 1] attention scores.
        ggml_tensor * kq = ggml_mul_mat(ctx, k, q_u);

        // Add relative position bias; soft_max_ext fuses the scale.
        kq                    = ggml_add(ctx, kq, matrix_bd);
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr, scale, /*max_bias=*/0.0f);

        // V transposed for the value matmul: [T_q, head_dim, n_head, 1].
        ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v, 1, 0, 2, 3));

        // attn_weights @ V: [head_dim, T_q, n_head, 1]
        o = ggml_mul_mat(ctx, v_t, kq_soft);

        // Merge heads: permute -> [head_dim, n_head, T_q] (contiguous for
        // the reshape) then reshape to [d_model, T_q].
        o = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
    }
    // Collapse heads back to d_model, keeping the batch at ne[2].
    o = ggml_reshape_3d(ctx, o, d_model, T_q, B);

    o = ggml_mul_mat(ctx, b.attn_out_w, o);
    if (b.attn_out_b != nullptr) {
        o = ggml_add(ctx, o, b.attn_out_b);
    }
    return o;
}

// ===========================================================================
// Conformer block (FF1 -> attn -> conv -> FF2 -> LN_out)
// ===========================================================================

ggml_tensor * build_conformer_block(ggml_context *        ctx,
                                    ggml_tensor *         x,
                                    ggml_tensor *         pos_emb,
                                    const BlockView &     b,
                                    const BlockParams &   params,
                                    const BlockObserver * obs) {
    // Observer dispatch helper (no-op when obs / its callback is null).
    const auto notify = [&](const char * tag, ggml_tensor * t) {
        if (obs != nullptr && obs->on_point != nullptr) {
            obs->on_point(obs->user, tag, t);
        }
    };

    // Macaron FF1: x = x + 0.5 * FF1(LN_ff1(x))
    x = macaron_ff_residual(ctx, x, b.norm_ff1_w, b.norm_ff1_b, b.ff1_lin1_w, b.ff1_lin1_b, b.ff1_lin2_w, b.ff1_lin2_b);
    notify("after_ff1", x);

    // Self-attention with relative position. Full residual.
    //
    // Streaming carry (params.streaming_channel_in != nullptr):
    //   1. new_cache = concat(prev_cache[T_q_new:], x_norm), cpy'd into
    //      streaming_channel_out (NeMo q_keep_size = T_q_new, cache stays at
    //      T_cache). KV-cache mode rotates last_k/last_v instead (kv_mode).
    //   2. Virtualize x_norm for K/V: concat(prev_cache, x_norm) so
    //      rel_pos_mhsa attends over T_virtual = T_cache + T_q_new keys.
    //      Queries are sliced to T_q_new (x_q_new), so the output already
    //      has T_q_new rows; pos_emb / chunked_mask are sized for the
    //      rectangular [T_virtual, T_q_new] geometry.
    // T_q_new comes from the running x's ne[1].
    {
        const bool    streaming = (params.streaming_channel_in != nullptr);
        // KV-cache streaming: keys/values are cached pre-projected, so
        // the channel cache (and its concat/write) is skipped entirely
        // and x_norm stays at the new frames throughout.
        const bool    kv_mode = streaming && params.streaming_kv_k_in != nullptr && params.streaming_kv_v_in != nullptr;
        const int64_t T_q_new = streaming ? x->ne[1] : 0;
        const int64_t T_cache = streaming ? params.streaming_channel_in->ne[1] : 0;
        ggml_tensor * x_q_new = nullptr;
        ggml_tensor * x_norm  = layer_norm(ctx, x, b.norm_attn_w, b.norm_attn_b);

        // Streaming cache rotation: new_cache = concat(prev[T_q_new:],
        // fresh) for the common T_q_new < cache-size case, else the
        // last cache-size rows of fresh. Shared by the channel cache
        // and the KV cache (identical rolling-window semantics).
        const auto emit_rotated_cache = [&](ggml_tensor * prev, ggml_tensor * fresh, ggml_tensor * out) {
            const int64_t Tc        = prev->ne[1];
            ggml_tensor * new_cache = nullptr;
            if (T_q_new < Tc) {
                ggml_tensor * prev_tail = ggml_view_2d(ctx, prev, prev->ne[0], Tc - T_q_new, prev->nb[1],
                                                       /*offset=*/T_q_new * prev->nb[1]);
                new_cache               = ggml_concat(ctx, prev_tail, fresh, /*dim=*/1);
            } else {
                new_cache = ggml_view_2d(ctx, fresh, fresh->ne[0], Tc, fresh->nb[1], (T_q_new - Tc) * fresh->nb[1]);
            }
            ggml_tensor * cpy = ggml_cpy(ctx, new_cache, out);
            ggml_build_forward_expand(params.streaming_graph, cpy);
        };

        if (streaming && !kv_mode) {
            // Emit the new cache slot (T_cache rows) before virtualizing
            // x_norm: prev_cache tail (T_cache - T_q_new rows) concatenated
            // with this chunk's x_norm, or x_norm's last T_cache rows when
            // T_q_new >= T_cache.
            if (params.streaming_channel_out != nullptr && params.streaming_graph != nullptr) {
                ggml_tensor * new_cache = nullptr;
                if (T_q_new < T_cache) {
                    ggml_tensor * prev_tail =
                        ggml_view_2d(ctx, params.streaming_channel_in, params.streaming_channel_in->ne[0],
                                     T_cache - T_q_new, params.streaming_channel_in->nb[1],
                                     /*offset=*/T_q_new * params.streaming_channel_in->nb[1]);
                    new_cache = ggml_concat(ctx, prev_tail, x_norm,
                                            /*dim=*/1);
                } else {
                    // x_norm's last T_cache rows.
                    new_cache = ggml_view_2d(ctx, x_norm, x_norm->ne[0], T_cache, x_norm->nb[1],
                                             (T_q_new - T_cache) * x_norm->nb[1]);
                }
                ggml_tensor * cpy = ggml_cpy(ctx, new_cache, params.streaming_channel_out);
                ggml_build_forward_expand(params.streaming_graph, cpy);
            }

            // Virtual-T attention with query slicing: the pre-concat
            // x_norm doubles as the query-only input, so attention output
            // is computed for just the T_q_new new frames while K/V span
            // the full virtual window (after the concat below).
            x_q_new = x_norm;
            x_norm  = ggml_concat(ctx, params.streaming_channel_in, x_norm, /*dim=*/1);
        }

        ggml_tensor * attn_out = nullptr;
        if (kv_mode) {
            // Project K/V for the new frames only, prepend the cached
            // window, and rotate the cache forward.
            ggml_tensor * k_new = ggml_mul_mat(ctx, b.attn_k_w, x_norm);
            if (b.attn_k_b != nullptr) {
                k_new = ggml_add(ctx, k_new, b.attn_k_b);
            }
            ggml_tensor * v_new = ggml_mul_mat(ctx, b.attn_v_w, x_norm);
            if (b.attn_v_b != nullptr) {
                v_new = ggml_add(ctx, v_new, b.attn_v_b);
            }

            ggml_tensor * k_all = ggml_concat(ctx, params.streaming_kv_k_in, k_new, /*dim=*/1);
            ggml_tensor * v_all = ggml_concat(ctx, params.streaming_kv_v_in, v_new, /*dim=*/1);

            if (params.streaming_kv_k_out != nullptr && params.streaming_kv_v_out != nullptr &&
                params.streaming_graph != nullptr) {
                emit_rotated_cache(params.streaming_kv_k_in, k_new, params.streaming_kv_k_out);
                emit_rotated_cache(params.streaming_kv_v_in, v_new, params.streaming_kv_v_out);
            }

            attn_out = rel_pos_mhsa(ctx, x_norm, pos_emb, b, params,
                                    /*x_q=*/nullptr, k_all, v_all);
        } else {
            // x_q_new is the new-frame query slice when streaming, or
            // nullptr offline (full self-attention). rel_pos_mhsa already
            // returns only the T_q_new query rows in the streaming case,
            // so no output slicing is needed.
            attn_out = rel_pos_mhsa(ctx, x_norm, pos_emb, b, params, x_q_new);
        }

        x = ggml_add(ctx, x, attn_out);
    }
    notify("after_attn", x);

    // Convolution module. Full residual.
    {
        ggml_tensor * x_norm   = layer_norm(ctx, x, b.norm_conv_w, b.norm_conv_b);
        ggml_tensor * conv_out = conv_module(ctx, x_norm, b, params);
        x                      = ggml_add(ctx, x, conv_out);
    }
    notify("after_conv", x);

    // Macaron FF2.
    x = macaron_ff_residual(ctx, x, b.norm_ff2_w, b.norm_ff2_b, b.ff2_lin1_w, b.ff2_lin1_b, b.ff2_lin2_w, b.ff2_lin2_b);
    notify("after_ff2", x);

    // Final per-block LayerNorm.
    x = layer_norm(ctx, x, b.norm_out_w, b.norm_out_b);
    notify("out", x);
    return x;
}

// Pre-encode (DwStridingSubsampling).

namespace {

// Name a tensor as `<prefix>.<suffix>` if prefix is non-null. Returns t
// unchanged.
ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, const char * suffix) {
    if (t == nullptr || prefix == nullptr) {
        return t;
    }
    char buf[64];
    std::snprintf(buf, sizeof(buf), "%s.%s", prefix, suffix);
    ggml_set_name(t, buf);
    return t;
}

}  // namespace

// Pre-encode subsampling stack. Op order matches NeMo's
// DwStridingSubsampling (conformer.py:206-328):
//
//   conv0 (standard, 1->C, k=3 s=2 p=1)  -> bias -> ReLU
//   conv2 (depthwise, C->C, groups=C)    -> bias -> conv3 (pointwise, C->C) -> bias -> ReLU
//   conv5 (depthwise, C->C, groups=C)    -> bias -> conv6 (pointwise, C->C) -> bias -> ReLU
//   permute + cont + reshape to flatten freq*channels into a single
//   feature axis matching `pre_encode_in = channels * (n_mels / 8)`
//   linear (out_w, out_b) -> [d_model, T_enc, 1, 1]
//
// Each conv's fused bias is an explicit ggml_add (no-op when null). The 2-D
// depthwise convs (conv2, conv5) honor policy.direct_dw_in_pre_encode.
ggml_tensor * build_pre_encode(ggml_context *        ctx,
                               const PreEncodeView & pe,
                               ggml_tensor *         mel_in,
                               const ConvPolicy &    policy,
                               const char *          name_prefix,
                               const char *          error_tag,
                               PreEncodeValidMasks * valid_masks) {
    // Masked-subsampling helper: zeros each conv intermediate's padded time
    // region after a ReLU. The mask is a graph input ne=[1, H_stage, 1, B]
    // (broadcasts over freq and channels); `slot` receives the handle for
    // the driver to fill. No-op when valid_masks is null.
    const int64_t pe_batch         = mel_in->ne[3];
    auto          apply_valid_mask = [&](ggml_tensor * x, ggml_tensor ** slot, const char * mname) -> ggml_tensor * {
        if (valid_masks == nullptr) {
            return x;
        }
        ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, x->ne[1], 1, pe_batch);
        ggml_set_name(m, mname);
        ggml_set_input(m);
        *slot = m;
        return ggml_mul(ctx, x, m);
    };

    // Transpose the mel input from [T_mel, n_mels, 1, 1] to
    // [n_mels, T_mel, 1, 1] = [W=F, H=T, IC, N], the NHWC order
    // ggml_conv_2d expects. The kernel is not spatially symmetric, so F/T
    // order matters — swapping them diverges the conv math.
    ggml_tensor * x = ggml_permute(ctx, mel_in,
                                   /*axis0=*/1, /*axis1=*/0,
                                   /*axis2=*/2, /*axis3=*/3);
    x               = ggml_cont(ctx, x);
    x               = name_prefixed(x, name_prefix, "mel_t");

    // Causal pre_encode: NeMo's CausalConv2D pads (left=k-1, right=stride-1)
    // on both spatial axes before the conv (p=0). Offline variants take the
    // op-side (k-1)/2 symmetric padding instead.
    const bool causal_pe  = policy.causal_pre_encode;
    const int  pe_p_op    = causal_pe ? 0 : 1;
    auto       pad_causal = [&](ggml_tensor * t) {
        if (!causal_pe) {
            return t;
        }
        // For k=3 / s=2: left=2, right=1 on both axes (zero F32 pads).
        const int left = 2, right = 1;
        auto      make_pad = [&](int width_w, int width_h) {
            ggml_tensor * p = ggml_new_tensor_4d(ctx, t->type, width_w > 0 ? width_w : t->ne[0],
                                                 width_h > 0 ? width_h : t->ne[1], t->ne[2], t->ne[3]);
            return ggml_fill(ctx, p, 0.0f);
        };
        // dim 0 (W = freq).
        ggml_tensor * pad_l = make_pad(left, /*h=*/0);
        t                   = ggml_concat(ctx, pad_l, t, /*dim=*/0);
        ggml_tensor * pad_r = make_pad(right, /*h=*/0);
        t                   = ggml_concat(ctx, t, pad_r, /*dim=*/0);
        // dim 1 (H = time). At this point ne[0] grew by left+right.
        ggml_tensor * pad_t = ggml_new_tensor_4d(ctx, t->type, t->ne[0], left, t->ne[2], t->ne[3]);
        pad_t               = ggml_fill(ctx, pad_t, 0.0f);
        t                   = ggml_concat(ctx, pad_t, t, /*dim=*/1);
        ggml_tensor * pad_b = ggml_new_tensor_4d(ctx, t->type, t->ne[0], right, t->ne[2], t->ne[3]);
        pad_b               = ggml_fill(ctx, pad_b, 0.0f);
        t                   = ggml_concat(ctx, t, pad_b, /*dim=*/1);
        return t;
    };

    // conv0 (standard 2D conv: 1 in, channels out, k=3 s=2)
    x = pad_causal(x);
    x = ggml_conv_2d(ctx, pe.conv0_w, x,
                     /*s0=*/2, /*s1=*/2,
                     /*p0=*/pe_p_op, /*p1=*/pe_p_op,
                     /*d0=*/1, /*d1=*/1);
    x = add_conv_bias(ctx, x, pe.conv0_b);
    x = name_prefixed(x, name_prefix, "conv0");
    x = ggml_relu(ctx, x);
    x = name_prefixed(x, name_prefix, "relu0");
    x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s1 : nullptr, "pre_encode.valid_mask.s1");

    // conv2 (depthwise: channels -> channels, groups=channels, k=3 s=2).
    // im2col path (conv_2d_dw_f32) when direct_dw_in_pre_encode is false.
    x = pad_causal(x);
    if (policy.direct_dw_in_pre_encode) {
        x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv2_w), x,
                                   /*s0=*/2, /*s1=*/2,
                                   /*p0=*/pe_p_op, /*p1=*/pe_p_op,
                                   /*d0=*/1, /*d1=*/1);
    } else {
        x = conv_2d_dw_f32(ctx, pe.conv2_w, x,
                           /*s0=*/2, /*s1=*/2,
                           /*p0=*/pe_p_op, /*p1=*/pe_p_op,
                           /*d0=*/1, /*d1=*/1);
    }
    x = add_conv_bias(ctx, x, pe.conv2_b);
    x = name_prefixed(x, name_prefix, "conv2");

    // conv3 (pointwise: channels -> channels, k=1 s=1 p=0)
    x = ggml_conv_2d(ctx, pe.conv3_w, x,
                     /*s0=*/1, /*s1=*/1,
                     /*p0=*/0, /*p1=*/0,
                     /*d0=*/1, /*d1=*/1);
    x = add_conv_bias(ctx, x, pe.conv3_b);
    x = name_prefixed(x, name_prefix, "conv3");
    x = ggml_relu(ctx, x);
    x = name_prefixed(x, name_prefix, "relu3");
    x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2");

    // conv5 (depthwise) -> conv6 (pointwise) -> ReLU
    x = pad_causal(x);
    if (policy.direct_dw_in_pre_encode) {
        x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv5_w), x,
                                   /*s0=*/2, /*s1=*/2,
                                   /*p0=*/pe_p_op, /*p1=*/pe_p_op,
                                   /*d0=*/1, /*d1=*/1);
    } else {
        x = conv_2d_dw_f32(ctx, pe.conv5_w, x,
                           /*s0=*/2, /*s1=*/2,
                           /*p0=*/pe_p_op, /*p1=*/pe_p_op,
                           /*d0=*/1, /*d1=*/1);
    }
    x = add_conv_bias(ctx, x, pe.conv5_b);
    x = name_prefixed(x, name_prefix, "conv5");

    x = ggml_conv_2d(ctx, pe.conv6_w, x,
                     /*s0=*/1, /*s1=*/1,
                     /*p0=*/0, /*p1=*/0,
                     /*d0=*/1, /*d1=*/1);
    x = add_conv_bias(ctx, x, pe.conv6_b);
    x = name_prefixed(x, name_prefix, "conv6");
    x = ggml_relu(ctx, x);
    x = name_prefixed(x, name_prefix, "relu6");
    x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3");

    // At this point ne = [F'=16, T_enc, channels=256, 1] where
    // T_enc = floor(T_mel / 8). Flatten (F', C) into one feature axis
    // matching pre_encode_in = channels * (n_mels / subsampling_factor).
    // The reference uses flat index = c*F' + f' (C slow, F' fast), which in
    // ggml ne is: permute F' to axis 0, C to axis 1, then collapse 0 and 1:
    //   permute (0, 2, 1, 3): [F', T, C, 1] -> [F', C, T, 1]
    //   ggml_cont; reshape -> [F'*C, T] = [pre_encode_in, T_enc]
    const int64_t F_prime       = x->ne[0];
    const int64_t T_enc         = x->ne[1];
    const int64_t C             = x->ne[2];
    // Utterance batch rides the conv "N" axis (ne[3]) through the stem; the
    // flatten + linear below land it on ne[2] == B for the conformer blocks.
    const int64_t B             = x->ne[3];
    const int64_t pre_encode_in = F_prime * C;

    // Sanity check: the linear layer expects exactly pre_encode_in features.
    if (pe.out_w != nullptr && pre_encode_in != pe.out_w->ne[0]) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                "%s encoder: pre_encode_in mismatch: "
                "F'*C=%lld but out_w expects %lld",
                error_tag != nullptr ? error_tag : "conformer", static_cast<long long>(pre_encode_in),
                static_cast<long long>(pe.out_w->ne[0]));
        return nullptr;
    }

    x = ggml_permute(ctx, x, /*axis0=*/0, /*axis1=*/2,
                     /*axis2=*/1, /*axis3=*/3);
    x = name_prefixed(x, name_prefix, "permuted");
    x = ggml_cont(ctx, x);
    x = name_prefixed(x, name_prefix, "flat_pre");
    // [F', C, T_enc, B] -> [F'*C, T_enc, B]. B at ne[2] after the collapse.
    x = ggml_reshape_3d(ctx, x, pre_encode_in, T_enc, B);
    x = name_prefixed(x, name_prefix, "flat");

    // Linear projection to d_model (out_w ne=[pre_encode_in, d_model]).
    x = ggml_mul_mat(ctx, pe.out_w, x);
    x = name_prefixed(x, name_prefix, "linear");

    // Add bias [d_model], broadcast to [d_model, T_enc, 1, 1].
    if (pe.out_b != nullptr) {
        const int64_t d_model = pe.out_b->ne[0];
        ggml_tensor * bias_4d = ggml_reshape_4d(ctx, pe.out_b, d_model, 1, 1, 1);
        x                     = ggml_add(ctx, x, bias_4d);
    }
    x = name_prefixed(x, name_prefix, "out");

    return x;
}

}  // namespace transcribe::conformer