transcribe-cpp-sys 0.0.10

// arch/medasr/model.cpp - MedASR (Google LASR-CTC) family handler.
//
// load() reads hparams, validates the tensor catalog, streams weight
// bytes into the backend buffer, fuses the per-block BatchNorm at the
// host level, and initializes the LASR mel frontend.
//
// init_context() allocates the per-context state.
//
// run() builds the encoder graph, uploads the mel (or the reference
// mel via TRANSCRIBE_MEDASR_MEL_FROM_REF), allocates and uploads the
// per-block fused-BN scale/bias inputs, computes the graph, reads
// back the CTC logits, performs greedy collapse (drop blanks, collapse
// repeats), and populates the result snapshot.
//
// run_batch is left nullptr — the central dispatcher falls back to
// calling run() once per utterance. A run_batch fast path can be wired
// later (Stage 4 Step 9) without changing this file's contract.

#include "medasr.h"

#include "encoder.h"
#include "weights.h"

#include "transcribe-arch.h"
#include "transcribe-batch-util.h"
#include "transcribe-debug.h"
#include "transcribe-load-common.h"
#include "transcribe-loader.h"
#include "transcribe-log.h"
#include "transcribe-meta.h"

#include "ggml.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "gguf.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <ios>
#include <memory>
#include <string>
#include <thread>
#include <type_traits>
#include <vector>

namespace transcribe::medasr {

extern const Arch arch;

static_assert(std::is_base_of_v<transcribe_model,   MedAsrModel>);
static_assert(std::is_base_of_v<transcribe_session, MedAsrSession>);

MedAsrSession::~MedAsrSession() {
    if (sched != nullptr) {
        ggml_backend_sched_free(sched);
        sched = nullptr;
    }
    if (compute_ctx != nullptr) {
        ggml_free(compute_ctx);
        compute_ctx = nullptr;
    }
    encoder_out = nullptr;
}

MedAsrModel::~MedAsrModel() {
    if (ctx_meta != nullptr) {
        ggml_free(ctx_meta);
        ctx_meta = nullptr;
    }
    if (backend_buffer != nullptr) {
        ggml_backend_buffer_free(backend_buffer);
        backend_buffer = nullptr;
    }
    for (auto it = plan.scheduler_list.rbegin();
         it != plan.scheduler_list.rend(); ++it)
    {
        ggml_backend_free(*it);
    }
    plan.scheduler_list.clear();
    plan.primary      = nullptr;
    plan.primary_kind = transcribe::BackendKind::Unknown;
}

namespace {

constexpr const char k_default_variant[] = "medasr";

// ---------------------------------------------------------------------------
// Input-length contract (see docs/input-limits.md). MedASR is a soft-window
// family: the encoder is a CTC Conformer with no hard context cap, but its
// RoPE was trained to enc_max_pos_emb encoder positions. Past that the rotary
// embeddings extrapolate and accuracy degrades, so over-window audio is WARNED
// and PROCESSED — never rejected, never altered. The decode numerics are
// untouched: this contract only observes the frame count and logs.
// ---------------------------------------------------------------------------

// Time downsampling from mel frames to encoder frames: each of the
// sub_n_layers subsampling convs strides by sub_stride (validated to 2x2 in
// read_medasr_hparams), so the effective factor is sub_stride ^ sub_n_layers
// (== 4 for the shipped checkpoint). Falls back to 4 if either hparam is
// missing so we never divide by an unset factor.
int medasr_subsample_factor(const MedAsrHParams & hp) {
    if (hp.enc_sub_stride <= 0 || hp.enc_sub_n_layers <= 0) {
        return 4;
    }
    int factor = 1;
    for (int i = 0; i < hp.enc_sub_n_layers; ++i) factor *= hp.enc_sub_stride;
    return factor;
}

// Milliseconds of 16 kHz audio per encoder frame: one mel frame is
// hop_length / sample_rate seconds, and the subsampling stack keeps one
// encoder frame per `subsample_factor` mel frames. For the shipped hparams
// (hop=160, sr=16000, 4x) this is 160 * 4 * 1000 / 16000 = 40 ms/frame
// (100 fps mel -> 25 fps encoder), matching the per-token stride used by the
// CTC decoder below. Returns 0 if a rate hparam is missing.
int64_t medasr_ms_per_encoder_frame(const MedAsrHParams & hp) {
    if (hp.fe_hop_length <= 0 || hp.fe_sample_rate <= 0) {
        return 0;
    }
    return static_cast<int64_t>(hp.fe_hop_length) *
           medasr_subsample_factor(hp) * 1000 / hp.fe_sample_rate;
}

// Advisory transcribe_capabilities::max_audio_ms: the RoPE-trained window,
// enc_max_pos_emb encoder frames * ms_per_encoder_frame. This is the point
// past which RoPE extrapolates; longer audio is accepted with a WARN (see
// run()), not rejected. Returns 0 ("unknown / unbounded") if enc_max_pos_emb
// or a frontend rate is missing, so a misconfigured model is never advertised
// with a wrong finite number. For the shipped hparams this is
// 10000 * 40 ms = 400000 ms (~400 s).
int64_t medasr_max_audio_ms(const MedAsrHParams & hp) {
    const int64_t ms_per_frame = medasr_ms_per_encoder_frame(hp);
    if (hp.enc_max_pos_emb <= 0 || ms_per_frame <= 0) {
        return 0;
    }
    return static_cast<int64_t>(hp.enc_max_pos_emb) * ms_per_frame;
}

transcribe_status load(Loader &                              loader,
                       const transcribe_model_load_params *  params,
                       transcribe_model **                   out_model)
{
    const int64_t t_load_start = ggml_time_us();

    auto m = std::make_unique<MedAsrModel>();
    m->arch      = &arch;
    m->t_load_us = 0;

    m->variant = loader.variant().empty() ? k_default_variant : loader.variant();
    m->backend.clear();

    apply_family_invariants(*m);
    m->caps.n_languages = 0;
    m->caps.languages   = nullptr;

    if (auto st = read_capability_kv(loader.gguf(), m->caps); st != TRANSCRIBE_OK)
        return st;
    if (auto st = read_languages_kv(loader.gguf(), *m); st != TRANSCRIBE_OK)
        return st;
    if (auto st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK)
        return st;
    if (auto st = read_medasr_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK)
        return st;

    // Publish the soft-window input-length advisory now that the RoPE range
    // and frontend rate are known (apply_family_invariants ran before the
    // hparams were read, so it could not set this). MedASR never rejects on
    // length; this is the window past which RoPE degrades. See
    // docs/input-limits.md.
    m->caps.max_audio_ms = medasr_max_audio_ms(m->hparams);

    // Reopen for tensor metadata.
    gguf_init_params init_params {};
    init_params.no_alloc = true;
    init_params.ctx      = &m->ctx_meta;
    gguf_context * gguf_data = gguf_init_from_file(loader.path().c_str(),
                                                   init_params);
    if (gguf_data == nullptr) return TRANSCRIBE_ERR_GGUF;

    if (auto st = build_medasr_weights(m->ctx_meta, m->hparams, m->weights);
        st != TRANSCRIBE_OK)
    {
        gguf_free(gguf_data);
        return st;
    }

    const transcribe_backend_request backend_req =
        (params != nullptr) ? params->backend : TRANSCRIBE_BACKEND_AUTO;
    if (auto st = transcribe::load_common::init_backends(
            backend_req, (params != nullptr) ? params->gpu_device : 0, "medasr", m->plan);
        st != TRANSCRIBE_OK)
    {
        gguf_free(gguf_data);
        return st;
    }
    m->backend = ggml_backend_name(m->plan.primary);
    m->primary_backend = m->plan.primary;

    ggml_backend_buffer_t weights_buffer =
        ggml_backend_alloc_ctx_tensors(m->ctx_meta, m->plan.primary);
    if (weights_buffer == nullptr) {
        gguf_free(gguf_data);
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr: ggml_backend_alloc_ctx_tensors failed");
        return TRANSCRIBE_ERR_GGUF;
    }
    m->backend_buffer = weights_buffer;
    ggml_backend_buffer_set_usage(weights_buffer,
                                  GGML_BACKEND_BUFFER_USAGE_WEIGHTS);

    if (auto st = transcribe::load_common::stream_tensor_data(
            loader.path(), gguf_data, m->ctx_meta, "medasr");
        st != TRANSCRIBE_OK)
    {
        gguf_free(gguf_data);
        return st;
    }

    // Fuse BatchNorm into host-resident scale/bias vectors. After this
    // point the raw bn.{weight,bias,running_mean,running_var} tensors
    // are no longer consulted.
    if (auto st = fuse_batch_norm(gguf_data, m->ctx_meta, m->hparams, m->weights);
        st != TRANSCRIBE_OK)
    {
        gguf_free(gguf_data);
        return st;
    }

    gguf_free(gguf_data);

    // LASR mel frontend (window + filterbank baked into the GGUF). The
    // shared MelFrontend handles n_fft=512 with Accelerate vDSP / cblas
    // primitives; medasr opts in via pad_mode="none" (center=False) and
    // log_clamp_min=1e-5 (natural log with floor-clamp, vs NeMo's
    // log(x+eps)). The GGUF filterbank is laid out [n_freq, n_mels]
    // row-major (ggml ne=[n_mels, n_freq]); MelConfig wants [n_mels,
    // n_freq] row-major, so we transpose once at load.
    {
        const MedAsrHParams & hp = m->hparams;
        const int n_freq = hp.fe_n_fft / 2 + 1;
        transcribe::MelConfig cfg {};
        cfg.sample_rate   = hp.fe_sample_rate;
        cfg.num_mels      = hp.fe_num_mels;
        cfg.n_fft         = hp.fe_n_fft;
        cfg.win_length    = hp.fe_win_length;
        cfg.hop_length    = hp.fe_hop_length;
        cfg.pre_emphasis  = 0.0f;
        cfg.f_min         = hp.fe_mel_lower_hz;
        cfg.f_max         = hp.fe_mel_upper_hz;
        cfg.pad_mode      = "none";
        cfg.window_type   = "hann_symmetric";
        cfg.normalize     = "none";
        cfg.log_clamp_min = hp.fe_log_clamp_min;

        if (m->weights.frontend_window != nullptr) {
            cfg.window.assign(hp.fe_win_length, 0.0f);
            ggml_backend_tensor_get(m->weights.frontend_window,
                                    cfg.window.data(), 0,
                                    cfg.window.size() * sizeof(float));
        }
        if (m->weights.frontend_mel_filterbank != nullptr) {
            std::vector<float> raw(
                static_cast<size_t>(hp.fe_num_mels) * n_freq);
            ggml_backend_tensor_get(m->weights.frontend_mel_filterbank,
                                    raw.data(), 0,
                                    raw.size() * sizeof(float));
            // raw is [n_freq, n_mels] row-major (mel fast). Transpose
            // into [n_mels, n_freq] row-major (freq fast) for MelConfig.
            cfg.filterbank.assign(
                static_cast<size_t>(hp.fe_num_mels) * n_freq, 0.0f);
            for (int k = 0; k < n_freq; ++k) {
                for (int mi = 0; mi < hp.fe_num_mels; ++mi) {
                    cfg.filterbank[static_cast<size_t>(mi) * n_freq + k] =
                        raw[static_cast<size_t>(k) * hp.fe_num_mels + mi];
                }
            }
        }
        m->mel.emplace(cfg);
    }

    m->t_load_us = ggml_time_us() - t_load_start;
    *out_model = m.release();
    return TRANSCRIBE_OK;
}

transcribe_status init_context(transcribe_model *                model,
                               const transcribe_session_params * params,
                               transcribe_session **             out_ctx)
{
    if (model->arch != &arch) return TRANSCRIBE_ERR_INVALID_ARG;

    auto gc = std::make_unique<MedAsrSession>();
    gc->model     = model;
    gc->n_threads = params->n_threads;
    gc->kv_type   = params->kv_type;

    *out_ctx = gc.release();
    return TRANSCRIBE_OK;
}

// Load reference mel sidecar from <dir>/mel.in.f32 (validate.py emits
// JSON; the .f32 sidecar is written next to it for fast bytes-only
// reload). For the encoder bring-up path we accept either filename.
transcribe_status load_ref_mel(const std::string & dir,
                               int                 n_mels,
                               int &               n_mel_frames,
                               std::vector<float> & out)
{
    auto try_open = [&](const std::string & path) -> std::ifstream {
        return std::ifstream(path, std::ios::binary | std::ios::ate);
    };
    std::string base = dir;
    if (!base.empty() && base.back() != '/') base += '/';

    std::ifstream f = try_open(base + "mel.in.f32");
    std::string used = base + "mel.in.f32";
    if (!f) {
        f = try_open(base + "frontend.mel.out.f32");
        used = base + "frontend.mel.out.f32";
    }
    if (!f) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr: cannot open mel ref under '%s' "
                     "(expected mel.in.f32 or frontend.mel.out.f32)",
                     dir.c_str());
        return TRANSCRIBE_ERR_FILE_NOT_FOUND;
    }

    const std::streamsize total_bytes = f.tellg();
    f.seekg(0, std::ios::beg);
    if (total_bytes <= 0 ||
        (static_cast<size_t>(total_bytes) % (sizeof(float) * n_mels)) != 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr: mel ref '%s' size %lld not divisible by "
                     "n_mels=%d * 4",
                     used.c_str(), static_cast<long long>(total_bytes),
                     n_mels);
        return TRANSCRIBE_ERR_GGUF;
    }
    const size_t n_elems = static_cast<size_t>(total_bytes) / sizeof(float);
    n_mel_frames = static_cast<int>(n_elems / n_mels);
    out.assign(n_elems, 0.0f);
    f.read(reinterpret_cast<char *>(out.data()),
           static_cast<std::streamsize>(total_bytes));
    if (static_cast<std::streamsize>(f.gcount()) != total_bytes) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "medasr: short read on mel ref");
        return TRANSCRIBE_ERR_GGUF;
    }
    return TRANSCRIBE_OK;
}

// Greedy CTC decode: argmax per frame, collapse repeats, drop blanks
// AND skip the LASR special tokens (<s>=1, </s>=2, <unk>=3) — matching
// the reference's `processor.batch_decode(skip_special_tokens=True)`.
// Without this, an utterance whose final CTC frame argmaxes to id 2
// (`</s>`) leaks the literal string `</s>` into the transcript, which
// scores as a +1 insertion against the reference text and lifts the
// dataset WER by ~0.2pp on test-clean. `<epsilon>` (id 0) is already
// dropped by the standard blank-skip; `<unk>` and `<s>` are explicit
// skips here to mirror the HF tokenizer's behaviour.
void decode_ctc_greedy(const float * logits,
                       int           vocab,
                       int           T_enc,
                       int           blank_id,
                       std::vector<int> & tokens,
                       std::vector<int> & frames)
{
    tokens.clear();
    frames.clear();
    int prev = -1;
    for (int t = 0; t < T_enc; ++t) {
        const float * row = logits + static_cast<size_t>(t) * vocab;
        int best = 0;
        float best_v = row[0];
        for (int v = 1; v < vocab; ++v) {
            if (row[v] > best_v) { best_v = row[v]; best = v; }
        }
        if (best == blank_id) { prev = -1; continue; }
        if (best == prev) continue;
        prev = best;
        // skip_special_tokens=True equivalent: drop ids 1 (<s>), 2 (</s>),
        // and 3 (<unk>). The blank id is checked above; this conditional
        // is unreachable when blank_id is itself in {1,2,3}, but the
        // explicit check keeps the intent local.
        if (best == 1 || best == 2 || best == 3) continue;
        tokens.push_back(best);
        frames.push_back(t);
    }
}

transcribe_status run(transcribe_session *      session,
                      const float *             pcm,
                      int                       n_samples,
                      const transcribe_run_params *)
{
    if (session == nullptr || pcm == nullptr || n_samples <= 0) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }
    auto * gc = static_cast<MedAsrSession *>(session);
    auto * gm = static_cast<MedAsrModel *>(gc->model);
    if (gm == nullptr || gm->plan.scheduler_list.empty()) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }
    if (gc->poll_abort()) return TRANSCRIBE_ERR_ABORTED;

    transcribe::debug::init();

    // ----- Mel acquisition -----------------------------------------------
    int mel_n_frames = 0;
    const int64_t t_mel_start = ggml_time_us();
    if (const char * ref_dir = std::getenv("TRANSCRIBE_MEDASR_MEL_FROM_REF");
        ref_dir != nullptr && ref_dir[0] != '\0')
    {
        if (auto st = load_ref_mel(ref_dir, gm->hparams.fe_num_mels,
                                   mel_n_frames, gc->mel_buf);
            st != TRANSCRIBE_OK) return st;
    } else {
        int out_n_mels = 0;
        std::vector<float> m_major;
        if (auto st = gm->mel->compute(pcm, static_cast<size_t>(n_samples),
                                       m_major, out_n_mels, mel_n_frames,
                                       gc->n_threads);
            st != TRANSCRIBE_OK) return st;
        // Shared MelFrontend writes [n_mels, T_mel] row-major (m * T + t);
        // medasr's encoder mel_in tensor (ne=[n_mels, T_mel, 1, B]) expects
        // T-major bytes (each frame is n_mels contiguous floats). One-shot
        // transpose into mel_buf.
        gc->mel_buf.assign(
            static_cast<size_t>(out_n_mels) * mel_n_frames, 0.0f);
        for (int mi = 0; mi < out_n_mels; ++mi) {
            const float * src = m_major.data() +
                                static_cast<size_t>(mi) * mel_n_frames;
            for (int t = 0; t < mel_n_frames; ++t) {
                gc->mel_buf[static_cast<size_t>(t) * out_n_mels + mi] = src[t];
            }
        }
    }
    gc->t_mel_us = ggml_time_us() - t_mel_start;
    if (mel_n_frames <= 0) return TRANSCRIBE_ERR_GGUF;

    // ----- Soft-window input-length advisory (see docs/input-limits.md) -----
    // MedASR has no hard context cap, but its RoPE was trained to
    // enc_max_pos_emb encoder frames. Predict the encoder frame count from the
    // mel frames via the two stride-2 subsampling convs (same formula as the
    // run_batch path's subsampling_t_out, applied twice). If it exceeds the
    // trained range, WARN once and PROCEED — we do not reject and do not touch
    // the numerics; accuracy may simply degrade past this point.
    {
        const int sub_k = gm->hparams.enc_sub_kernel;
        const int sub_s = gm->hparams.enc_sub_stride;
        auto sub_t = [&](int t_in) {
            return (t_in < sub_k) ? 0 : (t_in - sub_k) / sub_s + 1;
        };
        const int enc_frames = sub_t(sub_t(mel_n_frames));
        if (enc_frames > gm->hparams.enc_max_pos_emb) {
            const int64_t ms_per_frame =
                medasr_ms_per_encoder_frame(gm->hparams);
            const double audio_s =
                static_cast<double>(enc_frames) *
                static_cast<double>(ms_per_frame) / 1000.0;
            transcribe::log_msg(
                TRANSCRIBE_LOG_LEVEL_WARN,
                "medasr run: audio is %.1f s (%d encoder frames), beyond the "
                "%d-frame RoPE range this model was trained on; transcription "
                "may be degraded past this point. See "
                "transcribe_capabilities.max_audio_ms.",
                audio_s, enc_frames, gm->hparams.enc_max_pos_emb);
        }
    }

    // ----- Reset per-call compute state ---------------------------------
    if (gc->compute_ctx != nullptr) {
        ggml_free(gc->compute_ctx);
        gc->compute_ctx = nullptr;
    }
    gc->encoder_out = nullptr;
    {
        ggml_init_params ip {};
        ip.mem_size   = 4 * 1024 * 1024;
        ip.mem_buffer = nullptr;
        ip.no_alloc   = true;
        gc->compute_ctx = ggml_init(ip);
        if (gc->compute_ctx == nullptr) {
            transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                "medasr run: compute context allocation failed — out of memory. "
                "Split long audio into shorter segments (see "
                "transcribe_capabilities.max_audio_ms).");
            return TRANSCRIBE_ERR_OOM;
        }
    }

    // ----- Build encoder graph ------------------------------------------
    EncoderBuild eb = build_encoder_graph(
        gc->compute_ctx, gm->weights, gm->hparams,
        mel_n_frames, gm->backend.c_str(),
        /*n_batch=*/1, /*batch_var_len=*/false);
    if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) {
        return TRANSCRIBE_ERR_GGUF;
    }

    // ----- Scheduler ----------------------------------------------------
    if (gc->sched == nullptr) {
        gc->sched = ggml_backend_sched_new(
            gm->plan.scheduler_list.data(), nullptr,
            static_cast<int>(gm->plan.scheduler_list.size()),
            /*graph_size=*/8192, /*parallel=*/false, /*op_offload=*/true);
        if (gc->sched == nullptr) {
            transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                "medasr run: scheduler allocation failed — out of memory. "
                "Split long audio into shorter segments (see "
                "transcribe_capabilities.max_audio_ms).");
            return TRANSCRIBE_ERR_OOM;
        }
    }
    ggml_backend_sched_reset(gc->sched);
    if (!ggml_backend_sched_alloc_graph(gc->sched, eb.graph)) {
        transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
            "medasr run: encoder graph allocation failed — out of memory. "
            "Split long audio into shorter segments (see "
            "transcribe_capabilities.max_audio_ms).");
        return TRANSCRIBE_ERR_OOM;
    }

    // ----- Upload inputs ------------------------------------------------
    ggml_backend_tensor_set(eb.mel_in, gc->mel_buf.data(), 0,
                            gc->mel_buf.size() * sizeof(float));
    transcribe::debug::dump_tensor("mel.in", eb.mel_in, "encoder.mel");

    if (eb.positions != nullptr) {
        const int64_t T_enc = eb.positions->ne[0];
        std::vector<int32_t> pos(T_enc);
        for (int64_t i = 0; i < T_enc; ++i) pos[i] = static_cast<int32_t>(i);
        ggml_backend_tensor_set(eb.positions, pos.data(), 0,
                                pos.size() * sizeof(int32_t));
    }

    // Per-block fused-BN scale + bias.
    const int n_layers = gm->hparams.enc_n_layers;
    const int d_model  = gm->hparams.enc_hidden;
    const size_t bytes_per_layer = d_model * sizeof(float);
    for (int i = 0; i < n_layers; ++i) {
        if (i < static_cast<int>(eb.bn_scale_inputs.size()) &&
            eb.bn_scale_inputs[i] != nullptr) {
            ggml_backend_tensor_set(
                eb.bn_scale_inputs[i],
                gm->weights.fused_bn_scale_storage[i].data(),
                0, bytes_per_layer);
        }
        if (i < static_cast<int>(eb.bn_bias_inputs.size()) &&
            eb.bn_bias_inputs[i] != nullptr) {
            ggml_backend_tensor_set(
                eb.bn_bias_inputs[i],
                gm->weights.fused_bn_bias_storage[i].data(),
                0, bytes_per_layer);
        }
    }

    // ----- Compute ------------------------------------------------------
    const int64_t t_enc_start = ggml_time_us();
    if (ggml_backend_sched_graph_compute(gc->sched, eb.graph)
            != GGML_STATUS_SUCCESS) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "medasr: graph_compute failed");
        return TRANSCRIBE_ERR_GGUF;
    }
    gc->t_encode_us = ggml_time_us() - t_enc_start;

    // ----- Dump intermediates -------------------------------------------
    auto try_dump = [&](const char * name, ggml_tensor * t,
                        const char * stage) {
        if (t != nullptr) transcribe::debug::dump_tensor(name, t, stage);
    };
    try_dump("enc.subsampling.out", eb.dumps.subsampling_out, "encoder.subsampling");
    try_dump("enc.block.0.post_ff1",  eb.dumps.block0_post_ff1,  "encoder.block.0");
    try_dump("enc.block.0.post_attn", eb.dumps.block0_post_attn, "encoder.block.0");
    try_dump("enc.block.0.post_conv", eb.dumps.block0_post_conv, "encoder.block.0");
    try_dump("enc.block.0.post_ff2",  eb.dumps.block0_post_ff2,  "encoder.block.0");
    for (int i = 0; i < n_layers; ++i) {
        if (eb.dumps.all_block_outs[i] == nullptr) continue;
        char nm[32];
        std::snprintf(nm, sizeof(nm), "enc.block.%d.out", i);
        try_dump(nm, eb.dumps.all_block_outs[i], "encoder.block");
    }
    try_dump("enc.out_norm.out", eb.dumps.out_norm_out,        "encoder.out_norm");
    try_dump("enc.ctc_logits",   eb.dumps.ctc_logits_for_dump, "encoder.ctc_logits");

    // Diagnostic: per-block stats gated by TRANSCRIBE_MEDASR_DEBUG_STATS.
    // Prints min / max / mean / has_nan / has_inf for each marker tensor
    // to stderr (Modal captures and surfaces stderr in the run logs).
    // Used to localize where CUDA Q* output diverges from F32 / F16.
    if (const char * dbg = std::getenv("TRANSCRIBE_MEDASR_DEBUG_STATS");
        dbg != nullptr && dbg[0] != '\0')
    {
        std::vector<float> tmp;
        auto report = [&](const char * name, ggml_tensor * t) {
            if (t == nullptr) return;
            const size_t n = static_cast<size_t>(ggml_nelements(t));
            tmp.assign(n, 0.0f);
            ggml_backend_tensor_get(t, tmp.data(), 0, n * sizeof(float));
            float mn = +1e30f, mx = -1e30f;
            double sum = 0.0;
            int n_nan = 0, n_inf = 0;
            for (size_t i = 0; i < n; ++i) {
                const float v = tmp[i];
                if (std::isnan(v)) { ++n_nan; continue; }
                if (std::isinf(v)) { ++n_inf; continue; }
                sum += v;
                if (v < mn) mn = v;
                if (v > mx) mx = v;
            }
            const double mean = n > 0 ? sum / static_cast<double>(n) : 0.0;
            log_msg(TRANSCRIBE_LOG_LEVEL_DEBUG,
                         "STATS %-26s n=%zu  min=%+.4e  max=%+.4e  mean=%+.4e  nan=%d  inf=%d",
                         name, n, mn, mx, mean, n_nan, n_inf);
        };
        report("mel.in",                eb.mel_in);
        report("sub.after_dense0",      eb.dumps.sub_after_dense0);
        report("sub.after_conv0",       eb.dumps.sub_after_conv0);
        report("sub.after_conv1",       eb.dumps.sub_after_conv1);
        report("enc.subsampling.out",   eb.dumps.subsampling_out);
        report("enc.block.0.post_ff1",  eb.dumps.block0_post_ff1);
        report("enc.block.0.post_attn", eb.dumps.block0_post_attn);
        report("enc.block.0.post_conv", eb.dumps.block0_post_conv);
        report("enc.block.0.post_ff2",  eb.dumps.block0_post_ff2);
        for (int i = 0; i < n_layers; ++i) {
            if (eb.dumps.all_block_outs[i] == nullptr) continue;
            char nm[32];
            std::snprintf(nm, sizeof(nm), "enc.block.%d.out", i);
            report(nm, eb.dumps.all_block_outs[i]);
        }
        report("enc.out_norm.out", eb.dumps.out_norm_out);
        report("enc.ctc_logits",   eb.dumps.ctc_logits_for_dump);
    }

    // Per-utterance encoder-output dump (single-shot baseline for the
    // batched-equals-single-shot gate). Matches the parakeet / gigaam name.
    if (transcribe::debug::enabled() && eb.dumps.out_norm_out != nullptr) {
        ggml_tensor * t = eb.dumps.out_norm_out;
        const int d_enc = static_cast<int>(t->ne[0]);
        const int T_enc_dump = static_cast<int>(t->ne[1]);
        gc->enc_host.assign(static_cast<size_t>(d_enc) * T_enc_dump, 0.0f);
        ggml_backend_tensor_get(t, gc->enc_host.data(), 0,
                                gc->enc_host.size() * sizeof(float));
        const long long shape[2] = {T_enc_dump, d_enc};
        transcribe::debug::dump_host_f32("dec.enc_out", gc->enc_host.data(),
                                         static_cast<long long>(gc->enc_host.size()),
                                         shape, 2, "decoder.enc_out");
    }

    gc->encoder_out = eb.out;

    // ----- Greedy CTC decode --------------------------------------------
    const int64_t t_dec_start = ggml_time_us();
    ggml_tensor * logits_t = eb.dumps.ctc_logits;
    if (logits_t == nullptr) return TRANSCRIBE_ERR_GGUF;

    const int vocab = static_cast<int>(logits_t->ne[0]);
    const int T_enc = static_cast<int>(logits_t->ne[1]);
    if (vocab != gm->hparams.ctc_vocab_size || T_enc <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "medasr: ctc_logits shape mismatch (vocab=%d T_enc=%d)",
                     vocab, T_enc);
        return TRANSCRIBE_ERR_GGUF;
    }
    gc->logits_buf.assign(static_cast<size_t>(vocab) * T_enc, 0.0f);
    ggml_backend_tensor_get(logits_t, gc->logits_buf.data(), 0,
                            gc->logits_buf.size() * sizeof(float));

    std::vector<int> tokens;
    std::vector<int> frames;
    decode_ctc_greedy(gc->logits_buf.data(), vocab, T_enc,
                      gm->hparams.ctc_blank_id, tokens, frames);
    gc->t_decode_us = ggml_time_us() - t_dec_start;

    std::string text = gm->tok.decode(tokens.data(),
                                      static_cast<int>(tokens.size()));
    if (!text.empty() && text.front() == ' ') {
        text.erase(text.begin());
    }
    gc->full_text   = text;
    gc->has_result  = true;
    gc->result_kind = TRANSCRIBE_TIMESTAMPS_TOKEN;

    // Tokens at 40 ms per frame (16 kHz / hop=160 = 100 fps; subsampling
    // 4x reduces to 25 fps = 40 ms).
    gc->tokens.clear();
    gc->tokens.reserve(tokens.size());
    for (size_t i = 0; i < tokens.size(); ++i) {
        transcribe_session::TokenEntry te {};
        te.id    = tokens[i];
        te.text  = gm->tok.token(tokens[i]);
        te.t0_ms = static_cast<int64_t>(frames[i]) * 40;
        te.t1_ms = te.t0_ms + 40;
        gc->tokens.push_back(te);
    }

    return TRANSCRIBE_OK;
}

} // namespace

// ---------------------------------------------------------------------------
// Offline batch (transcribe_run_batch)
// ---------------------------------------------------------------------------
//
// Batches B utterances through one Conformer encoder dispatch (the batch
// rides ne[3] of the mel tensor + every activation in between), then
// host-decodes each utterance's CTC slice. The mel front-end and greedy
// decode are identical to single-shot; only the encoder is fused.
//
// Same-length batches run mask-free (bit-identical to single-shot per
// utterance). Variable-length batches pad to T_max along time and apply
// per-utterance masks (attn_pad_mask + conv_pad_mask), keeping a real
// query from attending to another utterance's padded tail.

namespace {

// Per-utterance post-subsampling frame count.
//   T_mel -> floor((T_mel - sub_kernel) / sub_stride) + 1     (conv_0)
//         -> floor((T1 - sub_kernel) / sub_stride) + 1        (conv_1)
// With sub_kernel=5 and sub_stride=2, padding=0, two stride-2 convs.
int subsampling_t_out(int T_in, int sub_kernel, int sub_stride) {
    if (T_in < sub_kernel) return 0;
    return (T_in - sub_kernel) / sub_stride + 1;
}

// medasr pack: source [b] is [T_b, n_mels] row-major (medasr mel
// orientation); destination is contiguous [n_mels-fast, T_max-slow,
// 1, n] byte-equivalent to numpy [n, T_max, n_mels]. Per utterance b
// the slab is at offset b * n_mels * T_max; valid time [0, T_b) is
// a memcpy of T_b*n_mels floats, tail stays zero. Matches the encoder
// mel_in tensor ne = [n_mels, T_max, 1, n].
void pack_pad_time_major(std::vector<float> &                    dst,
                         const std::vector<std::vector<float>> & src,
                         const std::vector<int> &                lens,
                         int                                     n_mels,
                         int                                     T_max)
{
    const size_t n = src.size();
    dst.assign(static_cast<size_t>(n_mels) * T_max * n, 0.0f);
    for (size_t b = 0; b < n; ++b) {
        const size_t slab_off = b * n_mels * T_max;
        const int T_b = lens[static_cast<size_t>(b)];
        const size_t copy_floats = static_cast<size_t>(T_b) * n_mels;
        if (T_b > 0 && copy_floats <= src[b].size()) {
            std::memcpy(dst.data() + slab_off, src[b].data(),
                        copy_floats * sizeof(float));
        }
        // Tail (t >= T_b) stays zero from the assign above.
    }
}

transcribe_status decode_one_utterance(MedAsrSession * gc,
                                       MedAsrModel *   gm,
                                       const float *   enc_logits,
                                       int             T_enc_valid,
                                       int             vocab,
                                       int             /*utt_index*/)
{
    const int64_t t_dec_start = ggml_time_us();

    std::vector<int> tokens;
    std::vector<int> frames;
    decode_ctc_greedy(enc_logits, vocab, T_enc_valid,
                      gm->hparams.ctc_blank_id, tokens, frames);

    gc->t_decode_us = ggml_time_us() - t_dec_start;

    std::string text = gm->tok.decode(tokens.data(),
                                      static_cast<int>(tokens.size()));
    if (!text.empty() && text.front() == ' ') {
        text.erase(text.begin());
    }
    gc->full_text   = text;
    gc->has_result  = true;
    gc->result_kind = TRANSCRIBE_TIMESTAMPS_TOKEN;

    gc->tokens.clear();
    gc->tokens.reserve(tokens.size());
    for (size_t i = 0; i < tokens.size(); ++i) {
        transcribe_session::TokenEntry te {};
        te.id    = tokens[i];
        te.text  = gm->tok.token(tokens[i]);
        te.t0_ms = static_cast<int64_t>(frames[i]) * 40;
        te.t1_ms = te.t0_ms + 40;
        gc->tokens.push_back(te);
    }
    return TRANSCRIBE_OK;
}

transcribe_status run_batch_encode(MedAsrSession *                         gc,
                                   MedAsrModel *                           gm,
                                   const std::vector<std::vector<float>> & mels,
                                   const std::vector<int> &                nf,
                                   int                                     n_mels,
                                   int                                     T_max,
                                   int64_t                                 total_mel_us)
{
    const int n = static_cast<int>(mels.size());
    bool var_len = false;
    for (int b = 0; b < n; ++b) {
        if (nf[b] != T_max) { var_len = true; break; }
    }

    pack_pad_time_major(gc->mel_buf, mels, nf, n_mels, T_max);

    if (gc->compute_ctx != nullptr) {
        ggml_free(gc->compute_ctx);
        gc->compute_ctx = nullptr;
    }
    gc->encoder_out = nullptr;
    {
        ggml_init_params ip {};
        ip.mem_size   = 4 * 1024 * 1024;
        ip.mem_buffer = nullptr;
        ip.no_alloc   = true;
        gc->compute_ctx = ggml_init(ip);
        if (gc->compute_ctx == nullptr) {
            transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                "medasr run: compute context allocation failed — out of memory. "
                "Split long audio into shorter segments (see "
                "transcribe_capabilities.max_audio_ms).");
            return TRANSCRIBE_ERR_OOM;
        }
    }

    EncoderBuild eb = build_encoder_graph(
        gc->compute_ctx, gm->weights, gm->hparams, T_max,
        gm->backend.c_str(), /*n_batch=*/n, /*batch_var_len=*/var_len);
    if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) {
        return TRANSCRIBE_ERR_GGUF;
    }

    if (gc->sched == nullptr) {
        gc->sched = ggml_backend_sched_new(
            gm->plan.scheduler_list.data(), nullptr,
            static_cast<int>(gm->plan.scheduler_list.size()),
            /*graph_size=*/8192, /*parallel=*/false, /*op_offload=*/true);
        if (gc->sched == nullptr) {
            transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                "medasr run: scheduler allocation failed — out of memory. "
                "Split long audio into shorter segments (see "
                "transcribe_capabilities.max_audio_ms).");
            return TRANSCRIBE_ERR_OOM;
        }
    }
    ggml_backend_sched_reset(gc->sched);
    if (!ggml_backend_sched_alloc_graph(gc->sched, eb.graph)) {
        transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
            "medasr run: encoder graph allocation failed — out of memory. "
            "Split long audio into shorter segments (see "
            "transcribe_capabilities.max_audio_ms).");
        return TRANSCRIBE_ERR_OOM;
    }

    // Upload mel + positions.
    ggml_backend_tensor_set(eb.mel_in, gc->mel_buf.data(), 0,
                            gc->mel_buf.size() * sizeof(float));
    if (eb.positions != nullptr) {
        const int64_t T_enc = eb.positions->ne[0];
        std::vector<int32_t> pos(T_enc);
        for (int64_t i = 0; i < T_enc; ++i) pos[i] = static_cast<int32_t>(i);
        ggml_backend_tensor_set(eb.positions, pos.data(), 0,
                                pos.size() * sizeof(int32_t));
    }

    // Per-layer fused-BN scale + bias (1-D, broadcast over the batch).
    const int n_layers = gm->hparams.enc_n_layers;
    const int d_model  = gm->hparams.enc_hidden;
    const size_t bytes_per_layer = d_model * sizeof(float);
    for (int i = 0; i < n_layers; ++i) {
        if (i < static_cast<int>(eb.bn_scale_inputs.size()) &&
            eb.bn_scale_inputs[i] != nullptr) {
            ggml_backend_tensor_set(
                eb.bn_scale_inputs[i],
                gm->weights.fused_bn_scale_storage[i].data(),
                0, bytes_per_layer);
        }
        if (i < static_cast<int>(eb.bn_bias_inputs.size()) &&
            eb.bn_bias_inputs[i] != nullptr) {
            ggml_backend_tensor_set(
                eb.bn_bias_inputs[i],
                gm->weights.fused_bn_bias_storage[i].data(),
                0, bytes_per_layer);
        }
    }

    // Per-utterance encoder-frame count after the two stride-2 subsampling
    // convs.
    std::vector<int> real_tenc(static_cast<size_t>(n), 0);
    const int sub_k = gm->hparams.enc_sub_kernel;
    const int sub_s = gm->hparams.enc_sub_stride;
    for (int b = 0; b < n; ++b) {
        int t = subsampling_t_out(nf[b], sub_k, sub_s);
        t     = subsampling_t_out(t,     sub_k, sub_s);
        real_tenc[b] = t;
    }

    // ----- Soft-window input-length advisory (see docs/input-limits.md) -----
    // Same warn-and-proceed contract as single-shot run(), applied per
    // utterance on the fused batched encode path so over-window clips in a
    // batch are not silently degraded. real_tenc[b] is the encoder frame count
    // (the same double-subsampling derivation run() uses); compare it to the
    // RoPE-trained range and WARN once per over-window utterance. The serial
    // fallback in run_batch() re-enters run(), which warns there, so each clip
    // is warned exactly once. Never reject, never touch the numerics.
    {
        const int64_t ms_per_frame = medasr_ms_per_encoder_frame(gm->hparams);
        for (int b = 0; b < n; ++b) {
            if (real_tenc[b] > gm->hparams.enc_max_pos_emb) {
                const double audio_s =
                    static_cast<double>(real_tenc[b]) *
                    static_cast<double>(ms_per_frame) / 1000.0;
                transcribe::log_msg(
                    TRANSCRIBE_LOG_LEVEL_WARN,
                    "medasr run: utterance %d: audio is %.1f s (%d encoder "
                    "frames), beyond the %d-frame RoPE range this model was "
                    "trained on; transcription may be degraded past this point. "
                    "See transcribe_capabilities.max_audio_ms.",
                    b, audio_s, real_tenc[b], gm->hparams.enc_max_pos_emb);
            }
        }
    }

    if (var_len) {
        const int T_enc_max = real_tenc[0];
        int T_max_eff = T_enc_max;
        for (int b = 1; b < n; ++b) {
            if (real_tenc[b] > T_max_eff) T_max_eff = real_tenc[b];
        }
        transcribe::fill_keypad_mask(eb.attn_pad_mask_in, real_tenc, T_max_eff, n);
        transcribe::fill_valid_frame_mask(eb.conv_pad_mask_in, real_tenc, T_max_eff, n);
    }

    // Per-backend thread count (matching single-shot through the scheduler).
    transcribe::configure_sched_n_threads(gc->sched, gc->n_threads);

    const int64_t t_enc_start = ggml_time_us();
    if (ggml_backend_sched_graph_compute(gc->sched, eb.graph)
            != GGML_STATUS_SUCCESS) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "medasr run_batch: graph_compute failed");
        return TRANSCRIBE_ERR_GGUF;
    }
    gc->t_encode_us = ggml_time_us() - t_enc_start;

    // Per-utterance encoder-output dump for the batched-equals-single-shot
    // gate. Read the whole [d_enc, T_enc, 1, n] tensor once, then host-slice
    // and dump each utterance under `dec.enc_out.b{i}`.
    if (transcribe::debug::enabled() && eb.dumps.out_norm_out != nullptr) {
        ggml_tensor * t = eb.dumps.out_norm_out;
        const int d_enc = static_cast<int>(t->ne[0]);
        const int T_enc_dump = static_cast<int>(t->ne[1]);
        const size_t utt_n = static_cast<size_t>(d_enc) * T_enc_dump;
        gc->enc_host.assign(utt_n * static_cast<size_t>(n), 0.0f);
        ggml_backend_tensor_get(t, gc->enc_host.data(), 0,
                                gc->enc_host.size() * sizeof(float));
        const long long shape[2] = {T_enc_dump, d_enc};
        for (int b = 0; b < n; ++b) {
            const std::string nm = "dec.enc_out.b" + std::to_string(b);
            transcribe::debug::dump_host_f32(
                nm.c_str(), gc->enc_host.data() + b * utt_n,
                static_cast<long long>(utt_n), shape, 2, "decoder.enc_out");
        }
    }

    // The decode tensor is eb.dumps.ctc_logits (ne=[vocab, T_enc, 1, n]),
    // not the transposed dump tensor. Byte layout per utterance b:
    // slab b at offset b * T_enc * vocab; within the slab, frame t at
    // offset t * vocab, then per-frame argmax over `vocab` floats.
    ggml_tensor * logits_t = eb.dumps.ctc_logits;
    if (logits_t == nullptr) return TRANSCRIBE_ERR_GGUF;
    const int vocab = static_cast<int>(logits_t->ne[0]);
    const int T_enc = static_cast<int>(logits_t->ne[1]);
    if (vocab != gm->hparams.ctc_vocab_size || T_enc <= 0) {
        return TRANSCRIBE_ERR_GGUF;
    }

    const size_t utt_elems = static_cast<size_t>(vocab) * T_enc;
    gc->logits_buf.assign(utt_elems * static_cast<size_t>(n), 0.0f);
    ggml_backend_tensor_get(logits_t, gc->logits_buf.data(), 0,
                            gc->logits_buf.size() * sizeof(float));

    // Host-decode each utterance with shared encode + total mel amortized.
    return transcribe::decode_batch_slices(
        gc, n, gc->logits_buf.data(), utt_elems, gc->t_encode_us, total_mel_us,
        [&](int b, const float * enc_b) {
            // enc_b points at utterance b's full T_enc slab; the decode
            // loop reads `real_tenc[b]` valid frames and ignores the padded
            // tail (per-utterance valid-frame count from the
            // subsampling-output formula above).
            return decode_one_utterance(gc, gm, enc_b, real_tenc[b], vocab, b);
        });
}

transcribe_status run_batch(transcribe_session *          session,
                            const float * const *         pcm,
                            const int *                   n_samples,
                            int                           n,
                            const transcribe_run_params * params)
{
    if (session == nullptr || pcm == nullptr || n_samples == nullptr || n <= 0) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }
    auto * gc = static_cast<MedAsrSession *>(session);
    auto * gm = static_cast<MedAsrModel *>(gc->model);
    if (gm == nullptr || gm->plan.scheduler_list.empty()) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }
    transcribe::debug::init();

    // Per-utterance mel in parallel.
    const int n_mels = gm->hparams.fe_num_mels;
    std::vector<std::vector<float>> mels(static_cast<size_t>(n));
    std::vector<int>                nf(static_cast<size_t>(n), 0);
    const int64_t t_mel_start = ggml_time_us();
    const bool all_ok = transcribe::parallel_for_all(
        n, gc->n_threads, [&](int i) -> bool {
            if (pcm[i] == nullptr || n_samples[i] <= 0) return false;
            int this_frames = 0;
            int out_n_mels = 0;
            std::vector<float> m_major;
            const transcribe_status st = gm->mel->compute(
                pcm[i], static_cast<size_t>(n_samples[i]),
                m_major, out_n_mels, this_frames, 1);
            if (st != TRANSCRIBE_OK || this_frames <= 0) return false;
            // Transpose [n_mels, T_mel] -> T-major to match the
            // encoder's mel_in tensor layout. See run() for the
            // detailed comment.
            mels[i].assign(
                static_cast<size_t>(out_n_mels) * this_frames, 0.0f);
            for (int mi = 0; mi < out_n_mels; ++mi) {
                const float * src = m_major.data() +
                                    static_cast<size_t>(mi) * this_frames;
                for (int t = 0; t < this_frames; ++t) {
                    mels[i][static_cast<size_t>(t) * out_n_mels + mi] = src[t];
                }
            }
            nf[i] = this_frames;
            return true;
        });
    const int64_t total_mel_us = ggml_time_us() - t_mel_start;

    if (all_ok) {
        int T_max = 0;
        for (int i = 0; i < n; ++i) T_max = std::max(T_max, nf[i]);
        return run_batch_encode(gc, gm, mels, nf, n_mels, T_max, total_mel_us);
    }

    // Per-utterance fallback (also the malformed-input path).
    for (int i = 0; i < n; ++i) {
        if (gc->poll_abort()) return TRANSCRIBE_ERR_ABORTED;
        if (pcm[i] == nullptr || n_samples[i] <= 0) {
            transcribe_session::ResultSet rs;
            rs.status = TRANSCRIBE_ERR_INVALID_ARG;
            gc->batch_results.push_back(std::move(rs));
            continue;
        }
        gc->clear_result();
        const transcribe_status st = run(session, pcm[i], n_samples[i], params);
        gc->batch_results.push_back(gc->capture_result(st));
    }
    return TRANSCRIBE_OK;
}

} // namespace

const Arch arch = {
    /*.name             =*/ "medasr",
    /*.load             =*/ load,
    /*.init_context     =*/ init_context,
    /*.run              =*/ run,
    /*.run_batch        =*/ run_batch,
    /*.stream_validate  =*/ nullptr,
    /*.stream_begin     =*/ nullptr,
    /*.stream_feed      =*/ nullptr,
    /*.stream_finalize  =*/ nullptr,
    /*.stream_reset     =*/ nullptr,
    /*.accepts_ext_kind =*/ nullptr,
    /*.run_validate     =*/ nullptr,
};

} // namespace transcribe::medasr