transcribe-cpp-sys 0.0.1

// arch/medasr/encoder.cpp - MedASR (Google LASR-CTC) Conformer encoder
// graph builder.
//
// Forward shape:
//
//   mel.in       ne=[T_mel, n_mels=128, 1, B]
//     -> subsampling (custom: dense + 2x conv1d(s=2,p=0) + dense)
//                  ne=[d_model=512, T_enc, 1, B]   = enc.subsampling.out
//     -> 17 x conformer block {
//         FF1 macaron residual w=[1.5, 0.5]   (sub-step: enc.block.0.post_ff1)
//         RoPE self-attention, unscaled residual (post_attn)
//         conv module, scaled residual w=[2.0, 1.0]                 (post_conv)
//         FF2 macaron residual w=[1.5, 0.5]                         (post_ff2)
//         norm_out                                          = enc.block.i.out
//       }
//     -> enc.out_norm                                 = enc.out_norm.out
//     -> ctc head (Conv1d k=1 = mul_mat + bias)
//                  ne=[vocab=512, T_enc, 1, B]                = enc.ctc_logits
//
// Differences from src/arch/gigaam/encoder.cpp (the closest in-tree
// analog):
//
//   - Subsampling is 1-D (linear -> relu -> conv1d(s=2) -> relu ->
//     conv1d(s=2) -> relu -> linear), NOT 2x conv1d. The leading dense
//     reduces n_mels -> d_model before the convs touch the time axis.
//   - All encoder LayerNorms are bias=false (norm has scale w, no beta).
//   - Conv module uses BatchNorm — host-fused at load time, uploaded as
//     compute_ctx tensors per call. Asymmetric padding (15, 16) for
//     k=32 "same" depthwise.
//   - Scaled residuals: macaron weights [1.5, 0.5] on FF1/FF2, conv
//     weights [2.0, 1.0]. NOT the standard 0.5 macaron half-step.
//   - Attention uses RoPE on Q and K AFTER projection (standard
//     convention). GigaAM rotates before projection, which is rare.
//   - LayerNorm eps is 1e-6 (loaded from hparams), NOT the
//     conf::layer_norm hardcoded 1e-5; this file uses a local helper
//     `lasr_layer_norm` that takes eps as a runtime arg.

#include "encoder.h"

#include "weights.h"

#include "conformer/conformer.h"
#include "transcribe-debug.h"
#include "transcribe-log.h"

#include "ggml.h"

#include <cmath>
#include <cstdio>
#include <cstring>
#include <string>

namespace transcribe::medasr {

namespace {

namespace conf = transcribe::conformer;

// Weight-matmul forcing F32 accumulation for F16 weights. On CUDA,
// `ggml_mul_mat` with an F16 weight takes the cuBLAS COMPUTE_16F path
// whose accumulator saturates at ~6.5e4 — and medasr's scaled-residual
// stream pushes intermediate activations to ~2e6 between blocks (the
// macaron [1.5, 0.5] and conv [2.0, 1.0] amplification). F16 accum
// overflows to NaN, the CTC argmax collapses to blank/specials, and
// every transcript comes back empty. GGML_PREC_F32 forces cuBLAS to run
// the GEMM in F32 instead, matching CPU/Metal which always F32-accumulate.
// No-op on CPU/Metal; slight perf cost on CUDA. Skip on non-F16 weights
// (BF16 already COMPUTE_32F; quantized routes through MMQ). Mirrors the
// `mul_mat_f32acc` helper at src/causal_lm/causal_lm.cpp:40.
ggml_tensor * mul_mat_f32acc(ggml_context * ctx,
                             ggml_tensor *  w,
                             ggml_tensor *  x)
{
    ggml_tensor * y = ggml_mul_mat(ctx, w, x);
    if (w->type == GGML_TYPE_F16) {
        ggml_mul_mat_set_prec(y, GGML_PREC_F32);
    }
    return y;
}

// LayerNorm with affine scale only (no bias). LASR uses
// `nn.LayerNorm(hidden_size, layer_norm_eps, bias=False)` on every
// encoder LN; eps is loaded from the GGUF (1e-6 for medasr) and is
// distinct from the conformer helper's hardcoded 1e-5.
ggml_tensor * lasr_layer_norm(ggml_context * ctx,
                              ggml_tensor *  x,
                              ggml_tensor *  scale,
                              float          eps)
{
    x = ggml_norm(ctx, x, eps);
    if (scale != nullptr) {
        x = ggml_mul(ctx, x, scale);
    }
    return x;
}

// Two-linear feed forward with no biases: y = Linear2(SiLU(Linear1(x))).
// Matches LASR's `LasrEncoderFeedForward` (SiLU activation via hidden_act).
ggml_tensor * lasr_feed_forward(ggml_context * ctx,
                                ggml_tensor *  x,
                                ggml_tensor *  up_w,
                                ggml_tensor *  down_w)
{
    x = mul_mat_f32acc(ctx, up_w, x);
    x = ggml_silu(ctx, x);
    x = mul_mat_f32acc(ctx, down_w, x);
    return x;
}

// Scaled residual: out = w0 * residual + w1 * branch. Used for FF1/FF2
// macaron (w0=1.5, w1=0.5) and conv (w0=2.0, w1=1.0). For the standard
// w0=1.0 residual we just ggml_add directly.
ggml_tensor * scaled_residual(ggml_context * ctx,
                              ggml_tensor *  residual,
                              ggml_tensor *  branch,
                              float          w0,
                              float          w1)
{
    ggml_tensor * a = (w0 == 1.0f) ? residual : ggml_scale(ctx, residual, w0);
    ggml_tensor * b = (w1 == 1.0f) ? branch   : ggml_scale(ctx, branch,   w1);
    return ggml_add(ctx, a, b);
}

// ---------------------------------------------------------------------------
// Subsampling
// ---------------------------------------------------------------------------
//
// Reference (LasrEncoderSubsampling.forward):
//   h = relu(dense_0(x))                # [B, T_mel, d_model]
//   h = h.transpose(1, 2)               # [B, d_model, T_mel]
//   h = relu(conv_0(h))                 # [B, d_model, T1]    (k=5, s=2, p=0)
//   h = relu(conv_1(h))                 # [B, sub_c, T_enc]   (k=5, s=2, p=0)
//   h = h.transpose(1, 2)               # [B, T_enc, sub_c]
//   return dense_1(h)                   # [B, T_enc, d_model] (no relu)
//
// In ggml (channel-fast):
//   mel_in     ne=[T_mel, n_mels, 1, B]
//   dense_0    mul_mat in [n_mels, T_mel, B] orientation -> [d_model, T_mel, B]
//              ... but mel_in is [T_mel, n_mels, 1, B]; we permute to
//              [n_mels, T_mel, 1, B] first so mul_mat sees n_mels at ne[0].
//   relu, then conv_0 takes [d_model, T_mel, B] as data with kernel
//   ne=[k=5, d_model, d_model]; conv_1d_f32 produces [T1, d_model, B] —
//   the same transpose pattern gigaam uses.
//   Pre-relu, transpose back to [d_model, T_mel, B] for the same-conv shape.
//
// We avoid back-and-forth permutes by keeping the running tensor in
// "channel-fast" [C, T, B] layout for the dense layers (mul_mat) and
// "time-fast" [T, C, B] layout for the conv layers (conv_1d_f32). One
// permute on each boundary is enough.
ggml_tensor * build_subsampling(ggml_context *            ctx,
                                const MedAsrSubsampling & ss,
                                ggml_tensor *             mel_in,
                                const MedAsrHParams &     hp,
                                const conf::ConvPolicy &  /*policy*/,
                                EncoderDumps *            dumps = nullptr)
{
    const int sub_k = hp.enc_sub_kernel;

    // mel_in ne=[n_mels, T_mel, 1, B]. Squeeze the singleton ne[2] so
    // dense_0's mul_mat sees a clean [n_mels, T_mel, B].
    ggml_tensor * x = ggml_reshape_3d(ctx, mel_in,
                                      mel_in->ne[0], mel_in->ne[1], mel_in->ne[3]);

    // dense_0: Linear(n_mels -> d_model). mul_mat broadcasts over the
    // batch axis (ne[2]).
    x = mul_mat_f32acc(ctx, ss.dense0_w, x);
    if (ss.dense0_b != nullptr) {
        x = ggml_add(ctx, x, ss.dense0_b);
    }
    x = ggml_relu(ctx, x);
    if (dumps != nullptr) {
        dumps->sub_after_dense0 = x;
        transcribe::debug::mark_tensor_for_dump(x);
    }
    // x ne=[d_model, T_mel, B].

    // Transpose to [T_mel, d_model, B] for conv_1d_f32 (which expects
    // data ne=[T, C, B] with kernel ne=[k, C_in, C_out]).
    x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

    // conv_0: Conv1d(d_model -> d_model, k=5, s=2, p=0). Output T1.
    x = conf::conv_1d_f32(ctx, ss.conv0_w, x, /*stride=*/sub_k > 0 ? hp.enc_sub_stride : 2,
                          /*padding=*/0, /*dilation=*/1);
    if (ss.conv0_b != nullptr) {
        ggml_tensor * bias_r = ggml_reshape_3d(ctx, ss.conv0_b, 1, x->ne[1], 1);
        x = ggml_add(ctx, x, bias_r);
    }
    x = ggml_relu(ctx, x);
    if (dumps != nullptr) {
        dumps->sub_after_conv0 = x;
        transcribe::debug::mark_tensor_for_dump(x);
    }
    // x ne=[T1, d_model, B].

    // conv_1: Conv1d(d_model -> sub_channels, k=5, s=2, p=0). Output T_enc.
    x = conf::conv_1d_f32(ctx, ss.conv1_w, x, /*stride=*/hp.enc_sub_stride,
                          /*padding=*/0, /*dilation=*/1);
    if (ss.conv1_b != nullptr) {
        ggml_tensor * bias_r = ggml_reshape_3d(ctx, ss.conv1_b, 1, x->ne[1], 1);
        x = ggml_add(ctx, x, bias_r);
    }
    x = ggml_relu(ctx, x);
    if (dumps != nullptr) {
        dumps->sub_after_conv1 = x;
        transcribe::debug::mark_tensor_for_dump(x);
    }
    // x ne=[T_enc, sub_channels, B].

    // Transpose to [sub_channels, T_enc, B] so dense_1's mul_mat sees
    // sub_channels at ne[0].
    x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));

    // dense_1: Linear(sub_channels -> d_model). NO RELU after this dense.
    //
    // Q4_K MMQ on CUDA overflows to +Inf for this specific tensor when
    // the activation magnitudes reach the ~3e4 range produced by the
    // ReLU-conv-ReLU-conv-ReLU stack above (see the per-stage stats in
    // the family-doc CUDA-quant debug section). The input dim here is
    // exactly 256 — one Q4_K super-block — and that shape + large
    // activations trip a CUDA-specific Q4_K MMQ bug. Workaround: cast
    // the dense_1 weight to F32 at graph build time on quants. The
    // tensor is small (256*512 = 128 KiB at F32) and runs once per
    // call, so the cost is negligible. F16/F32 GGUFs hit the cast
    // branch as a no-op (the type check skips them).
    ggml_tensor * dense1_w = ss.dense1_w;
    if (dense1_w->type != GGML_TYPE_F32 && dense1_w->type != GGML_TYPE_F16) {
        dense1_w = ggml_cast(ctx, dense1_w, GGML_TYPE_F32);
    }
    x = mul_mat_f32acc(ctx, dense1_w, x);
    if (ss.dense1_b != nullptr) {
        x = ggml_add(ctx, x, ss.dense1_b);
    }
    // x ne=[d_model, T_enc, B].
    return x;
}

// ---------------------------------------------------------------------------
// Conv module (BatchNorm path, asymmetric pad for k=32 "same")
// ---------------------------------------------------------------------------
//
// Adapts the conformer::conv_module shape but uses host-fused BN
// scale/bias passed in as input tensors (one pair per call). PyTorch's
// `padding="same"` with even k=32 and stride=1 splits total pad k-1=31
// as (left=15, right=16). conformer::conv_module supports asymmetric pad
// natively via params.conv_context_{left,right}.
ggml_tensor * build_conv_module(ggml_context *           ctx,
                                ggml_tensor *            x,
                                const MedAsrBlock &      b,
                                ggml_tensor *            bn_scale,
                                ggml_tensor *            bn_bias,
                                int                      d_model,
                                int                      conv_kernel,
                                const conf::ConvPolicy & policy,
                                ggml_tensor *            conv_pad_mask)
{
    conf::BlockView bv {};
    bv.conv_pw1_w          = b.conv_pw1_w;
    bv.conv_pw1_b          = nullptr;   // bias=False on LASR conv layers
    bv.conv_dw_w           = b.conv_dw_w;
    bv.conv_dw_b           = nullptr;
    bv.conv_pw2_w          = b.conv_pw2_w;
    bv.conv_pw2_b          = nullptr;
    bv.conv_bn_fused_scale = bn_scale;
    bv.conv_bn_fused_bias  = bn_bias;

    conf::BlockParams bp {};
    bp.d_model            = d_model;
    bp.n_head             = 0;          // unused in conv_module
    bp.conv_kernel        = conv_kernel;
    bp.kv_type            = GGML_TYPE_F32;
    bp.use_flash          = false;      // irrelevant for conv_module
    bp.policy             = policy;
    bp.conv_context_left  = (conv_kernel - 1) / 2;       // 15 for k=32
    bp.conv_context_right = conv_kernel - 1 - bp.conv_context_left; // 16
    bp.conv_norm_type     = conf::BlockParams::ConvNormType::BatchNorm;
    bp.conv_pad_mask      = conv_pad_mask;

    return conf::conv_module(ctx, x, bv, bp);
}

// ---------------------------------------------------------------------------
// RoPE self-attention
// ---------------------------------------------------------------------------
//
// Standard convention (matches LasrEncoderAttention.forward):
//   q = Wq @ x; k = Wk @ x; v = Wv @ x      (all no-bias)
//   reshape to [head_dim, n_head, T, B]
//   apply RoPE to q and k (positions tensor)
//   permute to [head_dim, T, n_head, B]
//   SDPA(q, k, v, mask=None)                 (scaling = 1/sqrt(head_dim))
//   permute + reshape to [d_model, T, B]
//   y = Wo @ ...                             (no bias)
//
// The full attention path matches gigaam::build_rotary_attn — the only
// medasr-specific deltas are:
//   - LASR rotates AFTER Wq/Wk projection (standard); gigaam rotates
//     BEFORE projection.
//   - rope_theta is 10000.0 (from hparams), not 5000.0.
ggml_tensor * build_rope_attn(ggml_context *      ctx,
                              ggml_tensor *       x,
                              ggml_tensor *       positions,
                              const MedAsrBlock & b,
                              int                 d_model,
                              int                 n_head,
                              float               rope_theta,
                              int                 rope_max_pos,
                              bool                use_flash,
                              ggml_tensor *       attn_pad_mask = nullptr)
{
    const int     head_dim = d_model / n_head;
    const float   scale    = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int64_t T        = x->ne[1];
    const int64_t Bb       = x->ne[2];

    // Per-utterance key-padding mask is content/length dependent — force
    // the manual SDPA path so the additive mask lands cleanly on the
    // score matrix before soft_max_ext's internal scale.
    const bool flash = use_flash && (attn_pad_mask == nullptr);

    // Q/K/V projections (no bias). Outputs [d_model, T, B].
    ggml_tensor * q = mul_mat_f32acc(ctx, b.attn_q_w, x);
    ggml_tensor * k = mul_mat_f32acc(ctx, b.attn_k_w, x);
    ggml_tensor * v = mul_mat_f32acc(ctx, b.attn_v_w, x);

    // Reshape Q/K to [head_dim, n_head, T, B] for RoPE rotation. ggml_rope_ext
    // rotates ne[0] (= head_dim) using position lookups on ne[2] (= T).
    q = ggml_reshape_4d(ctx, q, head_dim, n_head, T, Bb);
    k = ggml_reshape_4d(ctx, k, head_dim, n_head, T, Bb);

    q = ggml_rope_ext(ctx, q, positions, /*c=*/nullptr,
                      /*n_dims=*/head_dim, GGML_ROPE_TYPE_NEOX,
                      /*n_ctx_orig=*/rope_max_pos, rope_theta,
                      /*freq_scale=*/1.0f, /*ext_factor=*/0.0f,
                      /*attn_factor=*/1.0f, /*beta_fast=*/32.0f,
                      /*beta_slow=*/1.0f);
    k = ggml_rope_ext(ctx, k, positions, /*c=*/nullptr,
                      /*n_dims=*/head_dim, GGML_ROPE_TYPE_NEOX,
                      /*n_ctx_orig=*/rope_max_pos, rope_theta,
                      /*freq_scale=*/1.0f, /*ext_factor=*/0.0f,
                      /*attn_factor=*/1.0f, /*beta_fast=*/32.0f,
                      /*beta_slow=*/1.0f);

    // Reshape V to [head_dim, n_head, T, B]. No rotation on V.
    v = ggml_reshape_4d(ctx, v, head_dim, n_head, T, Bb);

    // SDPA target layout: [head_dim, T, n_head, B] so the attention
    // matmuls batch over (head, utterance) along ne[2..3].
    auto to_attn = [&](ggml_tensor * t) -> ggml_tensor * {
        t = ggml_permute(ctx, t, 0, 2, 1, 3);
        return ggml_cont(ctx, t);
    };
    q = to_attn(q);
    k = to_attn(k);
    v = to_attn(v);

    ggml_tensor * o;
    if (flash) {
        o = ggml_flash_attn_ext(ctx, q, k, v, /*mask=*/nullptr,
                                scale, 0.0f, 0.0f);
        // ggml_flash_attn_ext returns [head_dim, n_head, T, B].
        o = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
    } else {
        ggml_tensor * kq = ggml_mul_mat(ctx, k, q); // [T_k, T_q, n_head, B]
        if (attn_pad_mask != nullptr) {
            kq = ggml_add(ctx, kq, attn_pad_mask);
        }
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr,
                                                  scale, 0.0f);
        ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v, 1, 0, 2, 3));
        o = ggml_mul_mat(ctx, v_t, kq_soft);
    }
    // o: [head_dim, T, n_head, B] -> [head_dim, n_head, T, B] -> [d_model, T, B]
    o = ggml_permute(ctx, o, 0, 2, 1, 3);
    o = ggml_cont(ctx, o);
    o = ggml_reshape_3d(ctx, o, d_model, T, Bb);

    // Output projection (no bias).
    o = mul_mat_f32acc(ctx, b.attn_o_w, o);
    return o;
}

// ---------------------------------------------------------------------------
// Single Conformer block
// ---------------------------------------------------------------------------
struct BlockObs {
    ggml_tensor * post_ff1  = nullptr;
    ggml_tensor * post_attn = nullptr;
    ggml_tensor * post_conv = nullptr;
    ggml_tensor * post_ff2  = nullptr;
};

ggml_tensor * build_block(ggml_context *           ctx,
                          ggml_tensor *            x,
                          ggml_tensor *            positions,
                          const MedAsrBlock &      b,
                          ggml_tensor *            bn_scale,
                          ggml_tensor *            bn_bias,
                          const MedAsrHParams &    hp,
                          const conf::ConvPolicy & policy,
                          bool                     use_flash,
                          ggml_tensor *            attn_pad_mask,
                          ggml_tensor *            conv_pad_mask,
                          BlockObs *               obs)
{
    const int d_model    = hp.enc_hidden;
    const int n_head     = hp.enc_n_heads;
    const int conv_k     = hp.enc_conv_kernel;
    const float ln_eps   = hp.enc_layer_norm_eps;
    const float rope_th  = hp.enc_rope_theta;
    const int rope_max   = hp.enc_max_pos_emb;
    const float ff_w0    = hp.enc_ff_resid_w0;   // 1.5
    const float ff_w1    = hp.enc_ff_resid_w1;   // 0.5
    const float conv_w0  = hp.enc_conv_resid_w0; // 2.0
    const float conv_w1  = hp.enc_conv_resid_w1; // 1.0

    // FF1: x = w0*x + w1*FF1(LN(x)).
    {
        ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_ff1_w, ln_eps);
        y = lasr_feed_forward(ctx, y, b.ff1_up_w, b.ff1_down_w);
        x = scaled_residual(ctx, x, y, ff_w0, ff_w1);
    }
    if (obs != nullptr) obs->post_ff1 = x;

    // Self-attention: x = x + Attn(LN(x)).
    {
        ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_attn_w, ln_eps);
        y = build_rope_attn(ctx, y, positions, b, d_model, n_head,
                            rope_th, rope_max, use_flash, attn_pad_mask);
        x = ggml_add(ctx, x, y);
    }
    if (obs != nullptr) obs->post_attn = x;

    // Conv: x = w0*x + w1*Conv(LN(x)).
    {
        ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_conv_w, ln_eps);
        y = build_conv_module(ctx, y, b, bn_scale, bn_bias,
                              d_model, conv_k, policy, conv_pad_mask);
        x = scaled_residual(ctx, x, y, conv_w0, conv_w1);
    }
    if (obs != nullptr) obs->post_conv = x;

    // FF2: x = w0*x + w1*FF2(LN(x)).
    {
        ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_ff2_w, ln_eps);
        y = lasr_feed_forward(ctx, y, b.ff2_up_w, b.ff2_down_w);
        x = scaled_residual(ctx, x, y, ff_w0, ff_w1);
    }
    if (obs != nullptr) obs->post_ff2 = x;

    // Final per-block LN (= block output).
    x = lasr_layer_norm(ctx, x, b.norm_post_w, ln_eps);
    return x;
}

} // namespace

EncoderBuild build_encoder_graph(ggml_context *        ctx,
                                 const MedAsrWeights & w,
                                 const MedAsrHParams & hp,
                                 int                   n_mel_frames,
                                 const char *          backend_name,
                                 int                   n_batch,
                                 bool                  batch_var_len)
{
    EncoderBuild eb {};
    if (ctx == nullptr || n_mel_frames <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr encoder: invalid arg "
                     "(ctx=%p, n_mel_frames=%d)",
                     static_cast<void *>(ctx), n_mel_frames);
        return eb;
    }
    if (n_batch < 1) n_batch = 1;
    const bool var_len_masks = batch_var_len && n_batch > 1;

    const int d_model = hp.enc_hidden;
    const int n_layers = hp.enc_n_layers;

    // Mel input. The LASR reference dumps `mel.in` as numpy [T_mel, n_mels]
    // row-major (input_features shape after batch-squeeze). For ggml ne to
    // be byte-equivalent we need fast=n_mels, slow=T_mel — i.e.
    // ne=[n_mels, T_mel, 1, B]. The C++ MelFrontend writes the same
    // [T_mel, n_mels] row-major layout so the same upload byte stream
    // works in both modes.
    eb.mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32,
                                   hp.fe_num_mels, n_mel_frames, 1, n_batch);
    if (eb.mel_in == nullptr) return eb;
    ggml_set_name(eb.mel_in, "mel.in");
    ggml_set_input(eb.mel_in);
    transcribe::debug::mark_tensor_for_dump(eb.mel_in);

    // Conv policy. medasr's depthwise k=32 is too large for the direct
    // 2-D depthwise on some Metal kernels; use im2col im2col_dw to stay safe
    // until empirically profiled. Pointwise is always direct mul_mat.
    conf::ConvPolicy policy {};
    policy.direct_pw               = conf::detect_direct_pw(backend_name);
    policy.direct_dw_in_block      = true;
    policy.direct_dw_in_pre_encode = false;
    policy.causal_pre_encode       = false;

    // Build subsampling (passes dumps for stage-by-stage diagnostic).
    ggml_tensor * x = build_subsampling(ctx, w.subsampling, eb.mel_in, hp,
                                        policy, &eb.dumps);
    if (x == nullptr) return eb;
    ggml_set_name(x, "enc.subsampling.out");
    eb.dumps.subsampling_out = x;
    transcribe::debug::mark_tensor_for_dump(x);

    if (n_layers <= 0) {
        eb.out = x;
        eb.graph = ggml_new_graph(ctx);
        ggml_build_forward_expand(eb.graph, x);
        return eb;
    }

    const int64_t T_enc = x->ne[1];

    // Positions tensor for RoPE.
    eb.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_enc);
    ggml_set_name(eb.positions, "enc.positions");
    ggml_set_input(eb.positions);

    // Variable-length batch masks.
    if (var_len_masks) {
        eb.attn_pad_mask_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32,
                                                 T_enc, 1, 1, n_batch);
        eb.conv_pad_mask_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32,
                                                 T_enc, 1, n_batch, 1);
        ggml_set_name(eb.attn_pad_mask_in, "enc.attn_pad_mask");
        ggml_set_name(eb.conv_pad_mask_in, "enc.conv_pad_mask");
        ggml_set_input(eb.attn_pad_mask_in);
        ggml_set_input(eb.conv_pad_mask_in);
    }

    // Per-layer fused-BN scale/bias inputs. The driver uploads from
    // MedAsrWeights::fused_bn_*_storage after sched_alloc_graph.
    eb.bn_scale_inputs.assign(n_layers, nullptr);
    eb.bn_bias_inputs.assign(n_layers, nullptr);
    for (int i = 0; i < n_layers; ++i) {
        ggml_tensor * s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, d_model);
        ggml_tensor * bt = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, d_model);
        if (s == nullptr || bt == nullptr) return eb;
        char nm[64];
        std::snprintf(nm, sizeof(nm), "enc.block.%d.conv_bn.scale", i);
        ggml_set_name(s, nm);
        ggml_set_input(s);
        std::snprintf(nm, sizeof(nm), "enc.block.%d.conv_bn.bias", i);
        ggml_set_name(bt, nm);
        ggml_set_input(bt);
        eb.bn_scale_inputs[i] = s;
        eb.bn_bias_inputs[i]  = bt;
    }

    // Flash attention defaults on. Forced off when the manual SDPA path
    // is required (variable-length key-padding mask, single-shot/same-
    // length batches keep flash on).
    const bool use_flash = (std::getenv("TRANSCRIBE_NO_FLASH") == nullptr);

    eb.dumps.all_block_outs.assign(n_layers, nullptr);

    for (int i = 0; i < n_layers; ++i) {
        BlockObs obs{};
        const bool with_obs = (i == 0);
        x = build_block(ctx, x, eb.positions, w.blocks[i],
                        eb.bn_scale_inputs[i], eb.bn_bias_inputs[i],
                        hp, policy, use_flash,
                        eb.attn_pad_mask_in, eb.conv_pad_mask_in,
                        with_obs ? &obs : nullptr);

        if (with_obs) {
            // Block 0 sub-steps.
            const auto mark = [&](ggml_tensor * t, const char * nm,
                                  ggml_tensor *& slot) {
                if (t == nullptr) return;
                ggml_set_name(t, nm);
                transcribe::debug::mark_tensor_for_dump(t);
                slot = t;
            };
            mark(obs.post_ff1,  "enc.block.0.post_ff1",  eb.dumps.block0_post_ff1);
            mark(obs.post_attn, "enc.block.0.post_attn", eb.dumps.block0_post_attn);
            mark(obs.post_conv, "enc.block.0.post_conv", eb.dumps.block0_post_conv);
            mark(obs.post_ff2,  "enc.block.0.post_ff2",  eb.dumps.block0_post_ff2);
        }

        char nm[32];
        std::snprintf(nm, sizeof(nm), "enc.block.%d.out", i);
        ggml_set_name(x, nm);
        eb.dumps.all_block_outs[i] = x;
        // Mark every block output for dump; the dump pass only emits files for
        // the names referenced by the manifest, so marking extras is cheap.
        transcribe::debug::mark_tensor_for_dump(x);
    }

    // Top-level out norm.
    x = lasr_layer_norm(ctx, x, w.enc_out_norm_w, hp.enc_layer_norm_eps);
    ggml_set_name(x, "enc.out_norm.out");
    eb.dumps.out_norm_out = x;
    transcribe::debug::mark_tensor_for_dump(x);

    // CTC head: Conv1d(d_model -> vocab, k=1) = mul_mat. The converter
    // emits the weight as a k=1 conv with ne=[1, d_model, vocab]; reshape
    // to [d_model, vocab] for mul_mat (output ne=[vocab, T_enc, B], which
    // is byte-equivalent to numpy [T_enc, vocab] — handy for the decode
    // loop's per-frame argmax).
    {
        ggml_tensor * proj_w = ggml_reshape_2d(ctx, w.ctc_head.proj_w,
                                               d_model, hp.ctc_vocab_size);
        x = mul_mat_f32acc(ctx, proj_w, x);
        if (w.ctc_head.proj_b != nullptr) {
            x = ggml_add(ctx, x, w.ctc_head.proj_b);
        }
    }
    ggml_set_name(x, "enc.ctc_logits.raw");
    eb.dumps.ctc_logits = x;

    // Transposed view for the reference-comparable dump (numpy [vocab,
    // T_enc] order). The decode loop never reads this; only the dump pass.
    {
        ggml_tensor * logits_t = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
        ggml_set_name(logits_t, "enc.ctc_logits");
        eb.dumps.ctc_logits_for_dump = logits_t;
        transcribe::debug::mark_tensor_for_dump(logits_t);
    }

    eb.out = x;
    ggml_set_output(eb.out);
    if (eb.dumps.ctc_logits_for_dump != nullptr) {
        ggml_set_output(eb.dumps.ctc_logits_for_dump);
    }

    // Generous node budget. 17 blocks * ~50 ops/block + subsampling +
    // CTC head ~ 1000 nodes; 8192 leaves ample headroom.
    eb.graph = ggml_new_graph_custom(ctx, /*size=*/8192, /*grads=*/false);
    if (eb.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr encoder: ggml_new_graph_custom failed");
        return eb;
    }
    ggml_build_forward_expand(eb.graph, eb.out);
    // The transposed ctc_logits dump tensor is a SIDE output (reads from
    // eb.out but eb.out doesn't read from it), so it needs an explicit
    // forward-expand or the scheduler will skip it.
    if (eb.dumps.ctc_logits_for_dump != nullptr) {
        ggml_build_forward_expand(eb.graph, eb.dumps.ctc_logits_for_dump);
    }
    return eb;
}

} // namespace transcribe::medasr