transcribe-cpp-sys 0.0.5

// arch/whisper/decoder.cpp - Whisper decoder graph builder (prefill).
//
// Mirrors the transformers reference (WhisperDecoderLayer):
//
//   x = embed_tokens(ids) + embed_positions(weight[0:seq_len])
//   for layer in decoder.layers:
//       x = x + self_attn(LN_self(x), causal_mask)
//       x = x + cross_attn(LN_cross(x), encoder_hidden)
//       x = x + ffn(LN_ffn(x))
//   x = layer_norm(x)
//   logits_raw = proj_out(x)              # tied: W = embed_tokens.weight
//   logits     = log_softmax(logits_raw)
//
// Whisper has NO post-embed LayerNorm (cohere does — that is the
// notable shape difference vs. cohere/decoder.cpp).
//
// q/v/out carry bias on both self and cross attention; k_proj has NO
// bias. The shared `mha_decoder` helper takes nullable bias slots and
// skips the add when null, matching encoder.cpp's `mha_encoder`.

#include "decoder.h"

#include "encoder.h"      // for transcribe::conformer namespace alias context
#include "weights.h"
#include "whisper.h"

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

#include "ggml.h"

#include <algorithm>
#include <cmath>
#include <cstdio>

namespace transcribe::whisper {

namespace {

namespace conf = transcribe::conformer;
using conf::named;
using conf::layer_norm;

// Standard multi-head attention used by both self- and cross-attention
// in the decoder. context==nullptr selects self-attention (Q/K/V from x);
// non-null context selects cross-attention (Q from x, K/V from context).
//
// x:        [d_model, T_q]
// context:  [d_model, T_k]   or nullptr for self-attention (T_k = T_q)
// mask:     additive mask (f16) shape [T_k, T_q] or nullptr
// Returns:  [d_model, T_q]
ggml_tensor * mha_decoder(ggml_context * ctx,
                          ggml_tensor *  x,
                          ggml_tensor *  context,
                          ggml_tensor *  mask,
                          ggml_tensor *  q_w, ggml_tensor * q_b,
                          ggml_tensor *  k_w,
                          ggml_tensor *  v_w, ggml_tensor * v_b,
                          ggml_tensor *  out_w, ggml_tensor * out_b,
                          int            n_heads,
                          int            d_model,
                          bool           use_flash)
{
    const int     head_dim = d_model / n_heads;
    const float   scale    = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int64_t T_q      = x->ne[1];

    ggml_tensor * source = (context != nullptr) ? context : x;
    const int64_t T_k    = source->ne[1];

    ggml_tensor * q = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) q = ggml_add(ctx, q, q_b);

    ggml_tensor * k = ggml_mul_mat(ctx, k_w, source);
    // k has no bias.

    ggml_tensor * v = ggml_mul_mat(ctx, v_w, source);
    if (v_b != nullptr) v = ggml_add(ctx, v, v_b);

    // [d_model, T] -> [head_dim, n_heads, T] -> permute to
    // [head_dim, T, n_heads] for both flash and manual paths.
    q = ggml_reshape_3d(ctx, q, head_dim, n_heads, T_q);
    q = ggml_permute(ctx, q, 0, 2, 1, 3);

    k = ggml_reshape_3d(ctx, k, head_dim, n_heads, T_k);
    k = ggml_permute(ctx, k, 0, 2, 1, 3);

    v = ggml_reshape_3d(ctx, v, head_dim, n_heads, T_k);
    v = ggml_permute(ctx, v, 0, 2, 1, 3);

    ggml_tensor * o;
    if (use_flash) {
        ggml_tensor * q_c = ggml_cont(ctx, q);
        ggml_tensor * k_c = ggml_cont(ctx, k);
        ggml_tensor * v_c = ggml_cont(ctx, v);
        o = ggml_flash_attn_ext(ctx, q_c, k_c, v_c, mask, scale, 0.0f, 0.0f);
        o = ggml_reshape_2d(ctx, o, d_model, T_q);
    } else {
        ggml_tensor * kq      = ggml_mul_mat(ctx, ggml_cont(ctx, k),
                                             ggml_cont(ctx, q));
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, 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 = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
        o = ggml_reshape_2d(ctx, o, d_model, T_q);
    }

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

// FFN: fc1 -> GELU -> fc2. Both projections carry bias.
ggml_tensor * ffn(ggml_context * ctx,
                  ggml_tensor *  x,
                  ggml_tensor *  fc1_w, ggml_tensor * fc1_b,
                  ggml_tensor *  fc2_w, ggml_tensor * fc2_b) {
    ggml_tensor * h = ggml_mul_mat(ctx, fc1_w, x);
    if (fc1_b != nullptr) h = ggml_add(ctx, h, fc1_b);
    // HF Whisper uses nn.functional.gelu (exact erf form); ggml_gelu
    // is the tanh approximation and drifts ~1e-4 per element.
    h = ggml_gelu_erf(ctx, h);
    ggml_tensor * o = ggml_mul_mat(ctx, fc2_w, h);
    if (fc2_b != nullptr) o = ggml_add(ctx, o, fc2_b);
    return o;
}

ggml_tensor * build_block(ggml_context *          ctx,
                          ggml_tensor *           x,
                          ggml_tensor *           encoder_hidden,
                          ggml_tensor *           causal_mask,
                          const WhisperDecBlock & b,
                          int                     n_heads,
                          int                     d_model,
                          bool                    use_flash)
{
    // Self-attention (causal).
    {
        ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, b.norm_self_b);
        y = mha_decoder(ctx, y, /*context=*/nullptr, causal_mask,
                        b.self_q_w, b.self_q_b,
                        b.self_k_w,
                        b.self_v_w, b.self_v_b,
                        b.self_out_w, b.self_out_b,
                        n_heads, d_model, use_flash);
        x = ggml_add(ctx, x, y);
    }
    // Cross-attention (no mask; queries decoder state against encoder).
    {
        ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, b.norm_cross_b);
        y = mha_decoder(ctx, y, encoder_hidden, /*mask=*/nullptr,
                        b.cross_q_w, b.cross_q_b,
                        b.cross_k_w,
                        b.cross_v_w, b.cross_v_b,
                        b.cross_out_w, b.cross_out_b,
                        n_heads, d_model, use_flash);
        x = ggml_add(ctx, x, y);
    }
    // FFN.
    {
        ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, b.norm_ffn_b);
        y = ffn(ctx, y,
                b.ffn_fc1_w, b.ffn_fc1_b,
                b.ffn_fc2_w, b.ffn_fc2_b);
        x = ggml_add(ctx, x, y);
    }
    return x;
}

// ---------------------------------------------------------------------------
// KV-cached helpers
// ---------------------------------------------------------------------------
//
// Cache layout (flat 1D tensors, one slot per layer, per role):
//
//   self_k / self_v   : [hidden * n_layer * n_ctx]
//       layer il offset:  il * n_ctx * hidden
//       position  n  slot within layer:
//                         il * n_ctx * hidden + n * hidden
//
//   cross_k / cross_v : [hidden * n_layer * T_enc]
//       layer il offset:  il * T_enc * hidden
//
// Shape follows cohere's layout exactly so the ggml patterns transfer
// unchanged; whisper's only quirk is k_proj has no bias on both self
// and cross (cohere has k_b which may also be null).

// Self-attention reading/writing the self KV cache.
//
// x:         [d_model, n_tokens] input for new tokens only
// mask:      [n_kv, n_tokens] f16 additive mask, or nullptr
// il:        layer index (for cache offset)
// n_past:    number of tokens already in cache
// n_tokens:  number of new tokens being processed
//
// Writes new K/V into the cache at [il, n_past..n_past+n_tokens),
// then reads the full [il, 0..n_past+n_tokens) window for attention.
// Returns [d_model, n_tokens].
ggml_tensor * mha_self_cached(ggml_context *   ctx,
                              ggml_cgraph *    gf,
                              ggml_tensor *    x,
                              WhisperKvCache & kv_cache,
                              ggml_tensor *    mask,
                              ggml_tensor *    q_w, ggml_tensor * q_b,
                              ggml_tensor *    k_w,
                              ggml_tensor *    v_w, ggml_tensor * v_b,
                              ggml_tensor *    out_w, ggml_tensor * out_b,
                              int              n_heads,
                              int              d_model,
                              int              il,
                              int              n_past,
                              int              n_tokens,
                              int              n_kv,
                              bool             use_flash)
{
    const int head_dim = d_model / n_heads;
    const float scale  = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int n_ctx    = kv_cache.n_ctx;
    // n_kv may be padded to a FA-friendly multiple beyond n_past +
    // n_tokens. Trailing slots are masked to -inf via the caller's
    // mask so they do not contribute. The cache write below still
    // touches only the [n_past, n_past+n_tokens) range.

    // Q / K / V projections for the *new* tokens only.
    ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b);

    ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x);
    // k has no bias on whisper.

    ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x);
    if (v_b != nullptr) Vcur = ggml_add(ctx, Vcur, v_b);

    // Q in [head_dim, n_tokens, n_heads] for flash / manual paths.
    ggml_tensor * Q = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, n_tokens);
    Q = ggml_permute(ctx, Q, 0, 2, 1, 3);

    // Write new K/V into the cache. Layout is [hidden, n_ctx] per layer,
    // stored as a flat 1D buffer. Layer il occupies
    //   offset = il * n_ctx * hidden  (in elements).
    {
        const size_t k_elem = ggml_element_size(kv_cache.self_k);
        const size_t v_elem = ggml_element_size(kv_cache.self_v);

        ggml_tensor * k_dst = ggml_view_1d(
            ctx, kv_cache.self_k,
            static_cast<int64_t>(n_tokens) * d_model,
            k_elem * static_cast<size_t>(
                static_cast<int64_t>(il) * n_ctx * d_model +
                static_cast<int64_t>(n_past) * d_model));

        ggml_tensor * v_dst = ggml_view_1d(
            ctx, kv_cache.self_v,
            static_cast<int64_t>(n_tokens) * d_model,
            v_elem * static_cast<size_t>(
                static_cast<int64_t>(il) * n_ctx * d_model +
                static_cast<int64_t>(n_past) * d_model));

        ggml_build_forward_expand(gf, ggml_cpy(ctx, Kcur, k_dst));
        ggml_build_forward_expand(gf, ggml_cpy(ctx, Vcur, v_dst));
    }

    // Read K / V as [head_dim, n_kv, n_heads] for attention.
    const size_t k_elem = ggml_element_size(kv_cache.self_k);
    ggml_tensor * K = ggml_view_3d(
        ctx, kv_cache.self_k,
        head_dim, n_kv, n_heads,
        k_elem * d_model,        // nb1: stride between positions (one full d_model row)
        k_elem * head_dim,       // nb2: stride between heads inside a position
        k_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model));

    const size_t v_elem = ggml_element_size(kv_cache.self_v);
    ggml_tensor * V = ggml_view_3d(
        ctx, kv_cache.self_v,
        head_dim, n_kv, n_heads,
        v_elem * d_model,
        v_elem * head_dim,
        v_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model));

    ggml_tensor * o;
    if (use_flash) {
        o = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.0f, 0.0f);
        o = ggml_reshape_2d(ctx, o, d_model, n_tokens);
    } else {
        ggml_tensor * kq      = ggml_mul_mat(ctx, K, Q);
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, 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 = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
        o = ggml_reshape_2d(ctx, o, d_model, n_tokens);
    }

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

// Self-attention for the static-topology step graph. KV cache writes
// go through ggml_set_rows at runtime row index `kv_idx`; reads span
// the full [0, max_n_kv) window with `mask` gating valid positions.
// n_tokens is implicitly 1.
ggml_tensor * mha_self_step(ggml_context *   ctx,
                            ggml_cgraph *    gf,
                            ggml_tensor *    x,
                            WhisperKvCache & kv_cache,
                            ggml_tensor *    mask,
                            ggml_tensor *    kv_idx,
                            ggml_tensor *    q_w, ggml_tensor * q_b,
                            ggml_tensor *    k_w,
                            ggml_tensor *    v_w, ggml_tensor * v_b,
                            ggml_tensor *    out_w, ggml_tensor * out_b,
                            int              n_heads,
                            int              d_model,
                            int              il,
                            int              max_n_kv,
                            bool             use_flash)
{
    const int head_dim = d_model / n_heads;
    const float scale  = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int n_ctx    = kv_cache.n_ctx;

    ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b);

    ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x);
    // k has no bias on whisper.

    ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x);
    if (v_b != nullptr) Vcur = ggml_add(ctx, Vcur, v_b);

    ggml_tensor * Q = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, /*n_tokens=*/1);
    Q = ggml_permute(ctx, Q, 0, 2, 1, 3);

    // KV cache write via ggml_set_rows. Per-layer view spans [d_model, n_ctx];
    // we write one [d_model, 1] row at the kv_idx position.
    {
        const size_t k_elem = ggml_element_size(kv_cache.self_k);
        const size_t v_elem = ggml_element_size(kv_cache.self_v);
        const size_t layer_off_k = k_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model);
        const size_t layer_off_v = v_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model);

        ggml_tensor * k_layer = ggml_view_2d(
            ctx, kv_cache.self_k,
            d_model, n_ctx,
            k_elem * d_model, layer_off_k);
        ggml_tensor * v_layer = ggml_view_2d(
            ctx, kv_cache.self_v,
            d_model, n_ctx,
            v_elem * d_model, layer_off_v);

        ggml_tensor * K_row = ggml_reshape_2d(ctx, Kcur, d_model, 1);
        ggml_tensor * V_row = ggml_reshape_2d(ctx, Vcur, d_model, 1);

        ggml_build_forward_expand(
            gf, ggml_set_rows(ctx, k_layer, K_row, kv_idx));
        ggml_build_forward_expand(
            gf, ggml_set_rows(ctx, v_layer, V_row, kv_idx));
    }

    // Static read across [0, max_n_kv).
    const size_t k_elem = ggml_element_size(kv_cache.self_k);
    ggml_tensor * K = ggml_view_3d(
        ctx, kv_cache.self_k,
        head_dim, max_n_kv, n_heads,
        k_elem * d_model,
        k_elem * head_dim,
        k_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model));

    const size_t v_elem = ggml_element_size(kv_cache.self_v);
    ggml_tensor * V = ggml_view_3d(
        ctx, kv_cache.self_v,
        head_dim, max_n_kv, n_heads,
        v_elem * d_model,
        v_elem * head_dim,
        v_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * n_ctx * d_model));

    ggml_tensor * o;
    if (use_flash) {
        o = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.0f, 0.0f);
        o = ggml_reshape_2d(ctx, o, d_model, 1);
    } else {
        ggml_tensor * kq      = ggml_mul_mat(ctx, K, Q);
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, 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 = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
        o = ggml_reshape_2d(ctx, o, d_model, 1);
    }

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

// Cross-attention reading the pre-populated cross KV cache.
//
// x:        [d_model, n_tokens] query input
// il:       layer index
// T_enc:    number of encoder frames
//
// Returns: [d_model, n_tokens].
ggml_tensor * mha_cross_cached(ggml_context *   ctx,
                               ggml_tensor *    x,
                               WhisperKvCache & kv_cache,
                               ggml_tensor *    mask,
                               ggml_tensor *    q_w, ggml_tensor * q_b,
                               ggml_tensor *    out_w, ggml_tensor * out_b,
                               int              n_heads,
                               int              d_model,
                               int              il,
                               int              T_enc,
                               bool             use_flash)
{
    const int head_dim = d_model / n_heads;
    const float scale  = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int64_t n_tokens = x->ne[1];

    // Cross-attn cache is laid out at T_enc_pad rows per layer; the
    // trailing T_enc_pad - T_enc rows hold zeros (cleared at init,
    // never written by build_cross_kv_graph). FA reads the padded
    // shape — the unmasked padded slots contribute a small dilution
    // but produce a kernel-friendly sequence dim.
    const int T_enc_pad = kv_cache.T_enc_pad > 0
                              ? kv_cache.T_enc_pad : T_enc;
    (void)T_enc;

    // Q from current input; K / V read from the cross cache.
    ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b);

    ggml_tensor * Q = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, n_tokens);
    Q = ggml_permute(ctx, Q, 0, 2, 1, 3);

    const size_t k_elem = ggml_element_size(kv_cache.cross_k);
    ggml_tensor * K = ggml_view_3d(
        ctx, kv_cache.cross_k,
        head_dim, T_enc_pad, n_heads,
        k_elem * d_model,
        k_elem * head_dim,
        k_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * T_enc_pad * d_model));

    const size_t v_elem = ggml_element_size(kv_cache.cross_v);
    ggml_tensor * V = ggml_view_3d(
        ctx, kv_cache.cross_v,
        head_dim, T_enc_pad, n_heads,
        v_elem * d_model,
        v_elem * head_dim,
        v_elem * static_cast<size_t>(
            static_cast<int64_t>(il) * T_enc_pad * d_model));

    ggml_tensor * o;
    if (use_flash) {
        o = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.0f, 0.0f);
        o = ggml_reshape_2d(ctx, o, d_model, n_tokens);
    } else {
        ggml_tensor * kq      = ggml_mul_mat(ctx, K, Q);
        ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, 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 = ggml_permute(ctx, o, 0, 2, 1, 3);
        o = ggml_cont(ctx, o);
        o = ggml_reshape_2d(ctx, o, d_model, n_tokens);
    }

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

} // namespace

DecoderBuild build_decoder_prefill_graph(ggml_context *         ctx,
                                         const WhisperWeights & w,
                                         const WhisperHParams & hp,
                                         int                    seq_len,
                                         int                    T_enc,
                                         bool                   use_flash)
{
    DecoderBuild db {};

    if (ctx == nullptr || seq_len <= 0 || T_enc <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder: invalid arg "
                     "(ctx=%p, seq_len=%d, T_enc=%d)",
                     static_cast<void *>(ctx), seq_len, T_enc);
        return db;
    }
    if (seq_len > hp.dec_max_target_positions) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder: seq_len=%d exceeds "
                     "max_target_positions=%d",
                     seq_len, hp.dec_max_target_positions);
        return db;
    }

    const int d_model = hp.dec_d_model;
    const int n_heads = hp.dec_n_heads;
    const int vocab   = hp.dec_vocab_size;

    // ---- Inputs ------------------------------------------------------
    db.token_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, seq_len);
    named(db.token_ids_in, "dec.token_ids");
    ggml_set_input(db.token_ids_in);

    db.encoder_out_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_model, T_enc);
    named(db.encoder_out_in, "dec.encoder_out");
    ggml_set_input(db.encoder_out_in);

    // Causal mask declared as F32 input; cast to F16 inside the graph
    // because flash_attn_ext / soft_max_ext both want f16.
    ggml_tensor * causal_mask = nullptr;
    if (seq_len > 1) {
        db.causal_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32,
                                               seq_len, seq_len);
        named(db.causal_mask_in, "dec.causal_mask");
        ggml_set_input(db.causal_mask_in);
        causal_mask = ggml_cast(ctx, db.causal_mask_in, GGML_TYPE_F16);
    }

    // ---- Token embedding --------------------------------------------
    // token_embd_w is [d_model, vocab]; ggml_get_rows with i32 indices
    // [seq_len] yields [d_model, seq_len].
    ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w,
                                          db.token_ids_in);
    named(tok_emb, "dec.token_emb");
    transcribe::debug::mark_tensor_for_dump(tok_emb);
    db.dumps.token_emb = tok_emb;

    // ---- Positional embedding ---------------------------------------
    // pos_emb_w is [d_model, max_target_positions]. The transformers
    // reference takes weight[past:past+seq_len]; for prefill past=0 so
    // a 2D view of the leading seq_len rows is exactly the same data
    // (and avoids needing an i32 input for position ids).
    ggml_tensor * pos_emb = w.dec_top.pos_emb_w;
    if (seq_len != hp.dec_max_target_positions) {
        pos_emb = ggml_view_2d(ctx, w.dec_top.pos_emb_w,
                               d_model, seq_len,
                               w.dec_top.pos_emb_w->nb[1], 0);
    }
    named(pos_emb, "dec.pos_emb");
    transcribe::debug::mark_tensor_for_dump(pos_emb);
    db.dumps.pos_emb = pos_emb;

    // ---- Embedding sum (no LayerNorm) -------------------------------
    ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb);
    named(x, "dec.embed_sum");
    transcribe::debug::mark_tensor_for_dump(x);
    db.dumps.embed_sum = x;

    // ---- Decoder blocks ---------------------------------------------
    const int n_blocks = static_cast<int>(w.dec_blocks.size());
    db.dumps.block_outs.reserve(static_cast<size_t>(n_blocks));
    for (int i = 0; i < n_blocks; ++i) {
        x = build_block(ctx, x, db.encoder_out_in, causal_mask,
                        w.dec_blocks[i], n_heads, d_model, use_flash);

        char bname[64];
        std::snprintf(bname, sizeof(bname), "dec.block.%d.out", i);
        named(x, bname);
        transcribe::debug::mark_tensor_for_dump(x);
        db.dumps.block_outs.push_back(x);
    }

    // ---- Final LayerNorm --------------------------------------------
    x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b);
    named(x, "dec.out_before_head");
    transcribe::debug::mark_tensor_for_dump(x);
    db.dumps.out_before_head = x;

    // ---- Logits head (tied weight, no bias) -------------------------
    // mul_mat(W, x) where W=[d_model, vocab] and x=[d_model, seq_len]
    // gives [vocab, seq_len].
    ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x);
    named(logits_raw, "dec.logits_raw");
    transcribe::debug::mark_tensor_for_dump(logits_raw);
    db.dumps.logits_raw = logits_raw;

    ggml_tensor * logits = ggml_log(ctx, ggml_soft_max(ctx, logits_raw));
    named(logits, "dec.logits");
    transcribe::debug::mark_tensor_for_dump(logits);
    db.dumps.logits = logits;

    db.out = logits;
    ggml_set_output(db.out);

    // 8192 leaves headroom for 4-block whisper-tiny (and for larger
    // variants up to whisper-large-v3 at 32 blocks; each block adds
    // ~20 ops with self+cross+ffn).
    db.graph = ggml_new_graph_custom(ctx, 8192, false);
    if (db.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder: ggml_new_graph_custom failed");
        return db;
    }
    ggml_build_forward_expand(db.graph, db.out);

    // Suppress unused-variable warning when vocab is only used in the
    // bounds check above (kept for documentation in the catalog).
    (void)vocab;

    return db;
}

// ---------------------------------------------------------------------------
// KV-cached path: cross-KV precompute + prompt/step decoder graph.
// ---------------------------------------------------------------------------

DecoderBuild build_cross_kv_graph(ggml_context *         ctx,
                                  const WhisperWeights & w,
                                  const WhisperHParams & hp,
                                  WhisperKvCache &       kv_cache,
                                  ggml_tensor *          encoder_out,
                                  int                    T_enc)
{
    DecoderBuild db {};

    if (ctx == nullptr || encoder_out == nullptr || T_enc <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper cross_kv: invalid arg "
                     "(ctx=%p, encoder_out=%p, T_enc=%d)",
                     static_cast<void *>(ctx),
                     static_cast<void *>(encoder_out), T_enc);
        return db;
    }

    const int d_model = hp.dec_d_model;

    // View the persistent encoder_out tensor into this compute_ctx.
    // No data is copied — the view shares storage with the source,
    // which the encoder graph populated on the same backend.
    db.encoder_out_in = ggml_view_2d(ctx, encoder_out,
                                     d_model, T_enc,
                                     encoder_out->nb[1], 0);
    named(db.encoder_out_in, "dec.encoder_out");

    db.graph = ggml_new_graph_custom(ctx, 4096, false);
    if (db.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper cross_kv: ggml_new_graph_custom failed");
        return db;
    }

    // Layer slots in the flat cross cache are spaced by T_enc_pad
    // (allocator size) but each layer only writes T_enc rows; the
    // trailing T_enc_pad - T_enc rows stay zero from buffer_clear.
    const int T_enc_pad = kv_cache.T_enc_pad > 0
                              ? kv_cache.T_enc_pad : T_enc;

    const int n_layers = static_cast<int>(w.dec_blocks.size());
    for (int il = 0; il < n_layers; ++il) {
        const auto & blk = w.dec_blocks[il];

        // Whisper: cross k has no bias; cross v has a bias.
        ggml_tensor * Kcross = ggml_mul_mat(ctx, blk.cross_k_w,
                                            db.encoder_out_in);
        ggml_tensor * Vcross = ggml_mul_mat(ctx, blk.cross_v_w,
                                            db.encoder_out_in);
        if (blk.cross_v_b != nullptr) {
            Vcross = ggml_add(ctx, Vcross, blk.cross_v_b);
        }

        // Write the [d_model, T_enc] projections into the flat cache
        // slot for this layer. Offset in elements = il * T_enc_pad *
        // d_model so the layer slots in the padded cache do not
        // overlap; the size of each write is T_enc * d_model.
        const size_t k_elem = ggml_element_size(kv_cache.cross_k);
        const size_t v_elem = ggml_element_size(kv_cache.cross_v);

        ggml_tensor * k_dst = ggml_view_1d(
            ctx, kv_cache.cross_k,
            static_cast<int64_t>(T_enc) * d_model,
            k_elem * static_cast<size_t>(
                static_cast<int64_t>(il) * T_enc_pad * d_model));

        ggml_tensor * v_dst = ggml_view_1d(
            ctx, kv_cache.cross_v,
            static_cast<int64_t>(T_enc) * d_model,
            v_elem * static_cast<size_t>(
                static_cast<int64_t>(il) * T_enc_pad * d_model));

        ggml_build_forward_expand(db.graph, ggml_cpy(ctx, Kcross, k_dst));
        ggml_build_forward_expand(db.graph, ggml_cpy(ctx, Vcross, v_dst));
    }

    return db;
}

DecoderBuild build_decoder_graph_kv(ggml_context *         ctx,
                                    const WhisperWeights & w,
                                    const WhisperHParams & hp,
                                    WhisperKvCache &       kv_cache,
                                    int                    n_tokens,
                                    int                    n_past,
                                    int                    T_enc,
                                    int                    kv_pad,
                                    bool                   skip_log_softmax,
                                    bool                   use_flash)
{
    DecoderBuild db {};

    if (ctx == nullptr || n_tokens <= 0 || T_enc <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder_kv: invalid arg "
                     "(ctx=%p, n_tokens=%d, T_enc=%d)",
                     static_cast<void *>(ctx), n_tokens, T_enc);
        return db;
    }
    const int n_kv_active = n_past + n_tokens;
    if (n_kv_active > kv_cache.n_ctx) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder_kv: n_kv=%d exceeds n_ctx=%d",
                     n_kv_active, kv_cache.n_ctx);
        return db;
    }

    // Pad the K/V view length up to a FA-friendly multiple. Trailing
    // slots in [n_kv_active, n_kv) hold either zeros (first use of a
    // freshly-cleared cache) or stale K/V values from prior tiers /
    // chunks; the mask zeroes them out so they do not contribute.
    // Clamp to kv_cache.n_ctx so we do not read past the allocation.
    const int kv_pad_eff = kv_pad > 0 ? kv_pad : 1;
    int n_kv = n_kv_active;
    if (kv_pad_eff > 1) {
        n_kv = std::max(kv_pad_eff,
                        static_cast<int>(GGML_PAD(n_kv_active, kv_pad_eff)));
        if (n_kv > kv_cache.n_ctx) n_kv = kv_cache.n_ctx;
    }
    const bool has_padding = n_kv > n_kv_active;

    const int d_model = hp.dec_d_model;
    const int n_heads = hp.dec_n_heads;

    // ---- Inputs -----------------------------------------------------
    db.token_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
    named(db.token_ids_in, "dec.token_ids");
    ggml_set_input(db.token_ids_in);

    // Position ids as an input so the step loop can upload a single
    // int32 (the absolute position) without rebuilding the graph
    // structure. For prompt pass (n_past=0, n_tokens=4) the caller
    // uploads [0, 1, 2, 3].
    ggml_tensor * pos_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
    named(pos_ids_in, "dec.pos_ids");
    ggml_set_input(pos_ids_in);

    // Self-attention mask. Two reasons we need one:
    //   (1) prompt pass (n_tokens > 1) needs a causal mask
    //   (2) padded step (kv_pad > 1) needs -inf at trailing slots
    // The single mask covers both: shape [n_kv, n_tokens] with -inf
    // for k >= n_kv_active (covers padding) and -inf for k > q +
    // n_past (covers causality when n_tokens > 1). Caller in model.cpp
    // populates the values.
    //
    // When n_tokens == 1 AND there is no padding, the mask can be
    // omitted (single token attends to the full real cache window).
    // Preserve that fast path so non-FA / non-Metal paths skip the
    // extra input tensor.
    ggml_tensor * causal_mask = nullptr;
    const bool need_mask = (n_tokens > 1) || has_padding;
    if (need_mask) {
        db.causal_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32,
                                               n_kv, n_tokens);
        named(db.causal_mask_in, "dec.causal_mask");
        ggml_set_input(db.causal_mask_in);
        causal_mask = ggml_cast(ctx, db.causal_mask_in, GGML_TYPE_F16);
    } else {
        GGML_ASSERT(n_tokens == 1 &&
                    "decoder_kv mask-null branch requires single-token step");
    }

    // Cross-attention mask. Only needed when the cross cache has
    // been allocated with padding (T_enc_pad > T_enc); the trailing
    // padded slots hold zero K/V and would otherwise dilute the
    // unmasked FA softmax. -inf at those positions zeros the
    // contribution and restores numerical equivalence to an
    // unpadded view.
    ggml_tensor * cross_mask = nullptr;
    const int T_enc_pad = kv_cache.T_enc_pad > 0
                              ? kv_cache.T_enc_pad : T_enc;
    if (T_enc_pad > T_enc) {
        db.cross_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32,
                                              T_enc_pad, n_tokens);
        named(db.cross_mask_in, "dec.cross_mask");
        ggml_set_input(db.cross_mask_in);
        cross_mask = ggml_cast(ctx, db.cross_mask_in, GGML_TYPE_F16);
    }

    // ---- Token embedding -------------------------------------------
    ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w,
                                          db.token_ids_in);
    named(tok_emb, "dec.token_emb");
    if (n_past == 0) {
        transcribe::debug::mark_tensor_for_dump(tok_emb);
        db.dumps.token_emb = tok_emb;
    }

    // ---- Positional embedding --------------------------------------
    ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_top.pos_emb_w,
                                          pos_ids_in);
    named(pos_emb, "dec.pos_emb");
    if (n_past == 0) {
        transcribe::debug::mark_tensor_for_dump(pos_emb);
        db.dumps.pos_emb = pos_emb;
    }

    // ---- Embed sum (no LayerNorm) ----------------------------------
    ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb);
    named(x, "dec.embed_sum");
    if (n_past == 0) {
        transcribe::debug::mark_tensor_for_dump(x);
        db.dumps.embed_sum = x;
    }

    // Pre-create the graph so mha_self_cached can use
    // ggml_build_forward_expand to wire the cache writes into it.
    db.graph = ggml_new_graph_custom(ctx, 8192, false);
    if (db.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper decoder_kv: ggml_new_graph_custom failed");
        return db;
    }

    // ---- Decoder blocks --------------------------------------------
    const int n_blocks = static_cast<int>(w.dec_blocks.size());
    db.dumps.block_outs.reserve(static_cast<size_t>(n_blocks));
    for (int i = 0; i < n_blocks; ++i) {
        const auto & b = w.dec_blocks[i];

        // Self-attention (pre-LN, causal via mask when prompt pass).
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, b.norm_self_b);
            y = mha_self_cached(
                ctx, db.graph, y, kv_cache, causal_mask,
                b.self_q_w, b.self_q_b,
                b.self_k_w,
                b.self_v_w, b.self_v_b,
                b.self_out_w, b.self_out_b,
                n_heads, d_model,
                i, n_past, n_tokens, n_kv, use_flash);
            x = ggml_add(ctx, x, y);
        }
        // Cross-attention (pre-LN, reads cross cache).
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, b.norm_cross_b);
            y = mha_cross_cached(
                ctx, y, kv_cache, cross_mask,
                b.cross_q_w, b.cross_q_b,
                b.cross_out_w, b.cross_out_b,
                n_heads, d_model, i, T_enc, use_flash);
            x = ggml_add(ctx, x, y);
        }
        // FFN.
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, b.norm_ffn_b);
            y = ffn(ctx, y,
                    b.ffn_fc1_w, b.ffn_fc1_b,
                    b.ffn_fc2_w, b.ffn_fc2_b);
            x = ggml_add(ctx, x, y);
        }

        char bname[64];
        std::snprintf(bname, sizeof(bname), "dec.block.%d.out", i);
        named(x, bname);
        if (n_past == 0) {
            transcribe::debug::mark_tensor_for_dump(x);
            db.dumps.block_outs.push_back(x);
        }
    }

    // ---- Final LayerNorm -------------------------------------------
    x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b);
    named(x, "dec.out_before_head");
    if (n_past == 0) {
        transcribe::debug::mark_tensor_for_dump(x);
        db.dumps.out_before_head = x;
    }

    // ---- Logits head (tied weight, no bias) ------------------------
    ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x);
    named(logits_raw, "dec.logits_raw");
    if (n_past == 0) {
        transcribe::debug::mark_tensor_for_dump(logits_raw);
        db.dumps.logits_raw = logits_raw;
    }

    ggml_tensor * logits = logits_raw;
    if (!skip_log_softmax) {
        logits = ggml_log(ctx, ggml_soft_max(ctx, logits_raw));
        named(logits, "dec.logits");
        if (n_past == 0) {
            transcribe::debug::mark_tensor_for_dump(logits);
            db.dumps.logits = logits;
        }
    }

    db.out = logits;
    ggml_set_output(db.out);
    ggml_build_forward_expand(db.graph, db.out);

    return db;
}

// ---------------------------------------------------------------------------
// Static-topology single-token step graph (GPU dispatch path).
// ---------------------------------------------------------------------------

StepBuild build_step_graph(ggml_context *         ctx,
                           const WhisperWeights & w,
                           const WhisperHParams & hp,
                           WhisperKvCache &       kv_cache,
                           int                    max_n_kv,
                           int                    T_enc,
                           bool                   use_flash)
{
    StepBuild sb {};
    sb.max_n_kv = max_n_kv;

    if (ctx == nullptr || max_n_kv <= 0 || T_enc <= 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper step: invalid arg (max_n_kv=%d, T_enc=%d)",
                     max_n_kv, T_enc);
        return sb;
    }
    if (max_n_kv > kv_cache.n_ctx) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "whisper step: max_n_kv=%d exceeds kv_cache.n_ctx=%d",
                     max_n_kv, kv_cache.n_ctx);
        return sb;
    }

    const int d_model = hp.dec_d_model;
    const int n_heads = hp.dec_n_heads;
    const int T_enc_pad = kv_cache.T_enc_pad > 0 ? kv_cache.T_enc_pad : T_enc;
    const bool need_cross_mask = (T_enc_pad > T_enc);

    sb.token_id_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
    ggml_set_name(sb.token_id_in, "step.token_id");
    ggml_set_input(sb.token_id_in);

    sb.pos_id_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
    ggml_set_name(sb.pos_id_in, "step.pos_id");
    ggml_set_input(sb.pos_id_in);

    sb.kv_idx_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1);
    ggml_set_name(sb.kv_idx_in, "step.kv_idx");
    ggml_set_input(sb.kv_idx_in);

    sb.mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, max_n_kv, 1);
    ggml_set_name(sb.mask_in, "step.mask");
    ggml_set_input(sb.mask_in);

    if (need_cross_mask) {
        sb.cross_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_enc_pad, 1);
        ggml_set_name(sb.cross_mask_in, "step.cross_mask");
        ggml_set_input(sb.cross_mask_in);
    }
    ggml_tensor * cross_mask_f16 = need_cross_mask
        ? ggml_cast(ctx, sb.cross_mask_in, GGML_TYPE_F16)
        : nullptr;

    sb.graph = ggml_new_graph_custom(ctx, 8192, false);
    if (sb.graph == nullptr) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper step: ggml_new_graph_custom failed");
        return sb;
    }

    // Token + positional embedding (no post-embed LayerNorm).
    ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w, sb.token_id_in);
    ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_top.pos_emb_w,    sb.pos_id_in);
    ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb);

    const int n_blocks = static_cast<int>(w.dec_blocks.size());
    for (int i = 0; i < n_blocks; ++i) {
        const auto & b = w.dec_blocks[i];

        // Self-attention (pre-LN, set_rows-based KV write).
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, b.norm_self_b);
            y = mha_self_step(
                ctx, sb.graph, y, kv_cache, sb.mask_in, sb.kv_idx_in,
                b.self_q_w, b.self_q_b,
                b.self_k_w,
                b.self_v_w, b.self_v_b,
                b.self_out_w, b.self_out_b,
                n_heads, d_model,
                i, max_n_kv, use_flash);
            x = ggml_add(ctx, x, y);
        }
        // Cross-attention (pre-LN, reads pre-populated cross cache).
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, b.norm_cross_b);
            y = mha_cross_cached(
                ctx, y, kv_cache, cross_mask_f16,
                b.cross_q_w, b.cross_q_b,
                b.cross_out_w, b.cross_out_b,
                n_heads, d_model, i, T_enc, use_flash);
            x = ggml_add(ctx, x, y);
        }
        // FFN (pre-LN, GELU).
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, b.norm_ffn_b);
            y = ffn(ctx, y,
                    b.ffn_fc1_w, b.ffn_fc1_b,
                    b.ffn_fc2_w, b.ffn_fc2_b);
            x = ggml_add(ctx, x, y);
        }
    }

    x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b);

    // Tied lm_head, no bias. Output: raw pre-softmax logits.
    ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x);
    ggml_set_name(logits, "step.logits_raw");
    ggml_set_output(logits);
    sb.logits_out = logits;
    ggml_build_forward_expand(sb.graph, sb.logits_out);

    return sb;
}

// ===========================================================================
// Offline batched decode (B utterances)
// ===========================================================================

namespace {

// Batched self-attention step. x: [d_model, B] (one token per utterance).
// Mirrors mha_self_step but threads the utterance batch on the trailing
// flash axis (ne[3]); KV slab per layer is [d_model, n_ctx, B]. Flash-only.
ggml_tensor * mha_self_step_batched(
    ggml_context * ctx, ggml_cgraph * gf,
    ggml_tensor * x, WhisperKvCache & kv_cache,
    ggml_tensor * mask, ggml_tensor * kv_idx,
    ggml_tensor * q_w, ggml_tensor * q_b,
    ggml_tensor * k_w,
    ggml_tensor * v_w, ggml_tensor * v_b,
    ggml_tensor * out_w, ggml_tensor * out_b,
    int n_heads, int d_model, int il, int max_n_kv, int B)
{
    const int head_dim = d_model / n_heads;
    const float scale  = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const int64_t n_ctx = kv_cache.n_ctx;
    const size_t k_elem = ggml_element_size(kv_cache.self_k);
    const size_t v_elem = ggml_element_size(kv_cache.self_v);

    ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b);
    ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x);   // k has no bias
    ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x);
    if (v_b != nullptr) Vcur = ggml_add(ctx, Vcur, v_b);

    ggml_tensor * Q = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, B);

    // KV write: per-layer slab [d_model, n_ctx, B]; one set_rows writes B
    // rows at B independent indices kv_idx[b].
    {
        const size_t layer_off_k =
            k_elem * static_cast<size_t>(il) * n_ctx * d_model * B;
        const size_t layer_off_v =
            v_elem * static_cast<size_t>(il) * n_ctx * d_model * B;
        ggml_tensor * k_layer = ggml_view_3d(ctx, kv_cache.self_k,
            d_model, n_ctx, B, k_elem * d_model, k_elem * d_model * n_ctx, layer_off_k);
        ggml_tensor * v_layer = ggml_view_3d(ctx, kv_cache.self_v,
            d_model, n_ctx, B, v_elem * d_model, v_elem * d_model * n_ctx, layer_off_v);
        ggml_tensor * K_row = ggml_reshape_3d(ctx, Kcur, d_model, 1, B);
        ggml_tensor * V_row = ggml_reshape_3d(ctx, Vcur, d_model, 1, B);
        ggml_build_forward_expand(gf, ggml_set_rows(ctx, k_layer, K_row, kv_idx));
        ggml_build_forward_expand(gf, ggml_set_rows(ctx, v_layer, V_row, kv_idx));
    }

    const size_t layer_off_k =
        k_elem * static_cast<size_t>(il) * n_ctx * d_model * B;
    const size_t layer_off_v =
        v_elem * static_cast<size_t>(il) * n_ctx * d_model * B;
    ggml_tensor * K = ggml_view_4d(ctx, kv_cache.self_k,
        head_dim, max_n_kv, n_heads, B,
        k_elem * d_model, k_elem * head_dim, k_elem * d_model * n_ctx, layer_off_k);
    ggml_tensor * V = ggml_view_4d(ctx, kv_cache.self_v,
        head_dim, max_n_kv, n_heads, B,
        v_elem * d_model, v_elem * head_dim, v_elem * d_model * n_ctx, layer_off_v);

    // Q [head_dim, n_heads, B] -> [head_dim, 1, n_heads, B] for flash.
    ggml_tensor * Q_att = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 3, 1));
    ggml_tensor * o = ggml_flash_attn_ext(ctx, Q_att, K, V, mask, scale, 0.0f, 0.0f);
    o = ggml_reshape_2d(ctx, o, d_model, B);

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

// Batched cross-attention step. x: [d_model, B]; reads cross slab
// [d_model, T_enc_max, B] per layer with a per-utterance cross-pad mask.
// Whisper cross: q/out carry bias; k/v already live in the cache.
ggml_tensor * mha_cross_step_batched(
    ggml_context * ctx, ggml_tensor * x, WhisperKvCache & kv_cache,
    ggml_tensor * cross_mask,
    ggml_tensor * q_w, ggml_tensor * q_b,
    ggml_tensor * out_w, ggml_tensor * out_b,
    int n_heads, int d_model, int il, int T_enc_max, int B)
{
    const int head_dim = d_model / n_heads;
    const float scale  = 1.0f / std::sqrt(static_cast<float>(head_dim));
    const size_t k_elem = ggml_element_size(kv_cache.cross_k);
    const size_t v_elem = ggml_element_size(kv_cache.cross_v);

    ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x);
    if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b);
    ggml_tensor * Q = ggml_reshape_3d(ctx, Qcur, head_dim, n_heads, B);
    ggml_tensor * Q_att = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 3, 1));

    const size_t layer_off_k =
        k_elem * static_cast<size_t>(il) * T_enc_max * d_model * B;
    const size_t layer_off_v =
        v_elem * static_cast<size_t>(il) * T_enc_max * d_model * B;
    ggml_tensor * K = ggml_view_4d(ctx, kv_cache.cross_k,
        head_dim, T_enc_max, n_heads, B,
        k_elem * d_model, k_elem * head_dim, k_elem * d_model * T_enc_max, layer_off_k);
    ggml_tensor * V = ggml_view_4d(ctx, kv_cache.cross_v,
        head_dim, T_enc_max, n_heads, B,
        v_elem * d_model, v_elem * head_dim, v_elem * d_model * T_enc_max, layer_off_v);

    ggml_tensor * o = ggml_flash_attn_ext(ctx, Q_att, K, V, cross_mask, scale, 0.0f, 0.0f);
    o = ggml_reshape_2d(ctx, o, d_model, B);

    o = ggml_mul_mat(ctx, out_w, o);
    if (out_b != nullptr) o = ggml_add(ctx, o, out_b);
    return o;
}

} // namespace

DecoderBuild build_cross_kv_graph_batched(ggml_context *         ctx,
                                          const WhisperWeights & w,
                                          const WhisperHParams & hp,
                                          WhisperKvCache &       kv_cache,
                                          int                    T_enc_max,
                                          int                    n_batch)
{
    DecoderBuild db {};
    if (ctx == nullptr || T_enc_max <= 0 || n_batch <= 0) return db;

    const int64_t d_model = hp.dec_d_model;
    const int     B       = n_batch;

    db.encoder_out_in = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_model, T_enc_max, B);
    ggml_set_name(db.encoder_out_in, "dec.encoder_out");
    ggml_set_input(db.encoder_out_in);

    db.graph = ggml_new_graph_custom(ctx, 8192, false);
    if (db.graph == nullptr) return db;

    const size_t k_elem = ggml_element_size(kv_cache.cross_k);
    const size_t v_elem = ggml_element_size(kv_cache.cross_v);

    const int n_layers = static_cast<int>(w.dec_blocks.size());
    for (int il = 0; il < n_layers; ++il) {
        const auto & blk = w.dec_blocks[il];

        // Whisper: cross k has no bias; cross v has a bias.
        ggml_tensor * Kc = ggml_mul_mat(ctx, blk.cross_k_w, db.encoder_out_in);
        ggml_tensor * Vc = ggml_mul_mat(ctx, blk.cross_v_w, db.encoder_out_in);
        if (blk.cross_v_b != nullptr) Vc = ggml_add(ctx, Vc, blk.cross_v_b);
        // Kc/Vc: [d_model, T_enc_max, B].

        const size_t off_k =
            k_elem * static_cast<size_t>(il) * T_enc_max * d_model * B;
        const size_t off_v =
            v_elem * static_cast<size_t>(il) * T_enc_max * d_model * B;
        ggml_tensor * k_dst = ggml_view_3d(ctx, kv_cache.cross_k,
            d_model, T_enc_max, B, k_elem * d_model, k_elem * d_model * T_enc_max, off_k);
        ggml_tensor * v_dst = ggml_view_3d(ctx, kv_cache.cross_v,
            d_model, T_enc_max, B, v_elem * d_model, v_elem * d_model * T_enc_max, off_v);
        ggml_build_forward_expand(db.graph, ggml_cpy(ctx, Kc, k_dst));
        ggml_build_forward_expand(db.graph, ggml_cpy(ctx, Vc, v_dst));
    }
    return db;
}

StepBuildBatched build_step_graph_batched(ggml_context *         ctx,
                                          const WhisperWeights & w,
                                          const WhisperHParams & hp,
                                          WhisperKvCache &       kv_cache,
                                          int                    max_n_kv,
                                          int                    T_enc_max,
                                          int                    n_batch,
                                          bool                   use_flash)
{
    StepBuildBatched sb {};
    sb.max_n_kv = max_n_kv;
    sb.n_batch  = n_batch;
    if (ctx == nullptr || max_n_kv <= 0 || T_enc_max <= 0 || n_batch <= 0) return sb;
    if (!use_flash) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper step(batched): requires flash path");
        return sb;
    }

    const int d_model = hp.dec_d_model;
    const int n_heads = hp.dec_n_heads;
    const int B       = n_batch;

    sb.token_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, B);
    ggml_set_name(sb.token_ids_in, "step.token_ids"); ggml_set_input(sb.token_ids_in);
    sb.pos_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, B);
    ggml_set_name(sb.pos_ids_in, "step.pos_ids"); ggml_set_input(sb.pos_ids_in);
    sb.kv_idx_in = ggml_new_tensor_2d(ctx, GGML_TYPE_I64, 1, B);
    ggml_set_name(sb.kv_idx_in, "step.kv_idx"); ggml_set_input(sb.kv_idx_in);
    sb.self_mask_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, max_n_kv, 1, 1, B);
    ggml_set_name(sb.self_mask_in, "step.self_mask"); ggml_set_input(sb.self_mask_in);
    sb.cross_mask_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, T_enc_max, 1, 1, B);
    ggml_set_name(sb.cross_mask_in, "step.cross_mask"); ggml_set_input(sb.cross_mask_in);

    sb.graph = ggml_new_graph_custom(ctx, 16384, false);
    if (sb.graph == nullptr) return sb;

    // Token + positional embedding (no post-embed LayerNorm).
    ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w, sb.token_ids_in);
    ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_top.pos_emb_w,    sb.pos_ids_in);
    ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb);

    const int n_blocks = static_cast<int>(w.dec_blocks.size());
    for (int i = 0; i < n_blocks; ++i) {
        const auto & b = w.dec_blocks[i];
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, b.norm_self_b);
            ggml_tensor * so = mha_self_step_batched(
                ctx, sb.graph, y, kv_cache, sb.self_mask_in, sb.kv_idx_in,
                b.self_q_w, b.self_q_b, b.self_k_w,
                b.self_v_w, b.self_v_b, b.self_out_w, b.self_out_b,
                n_heads, d_model, i, max_n_kv, B);
            x = ggml_add(ctx, x, so);
        }
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, b.norm_cross_b);
            ggml_tensor * co = mha_cross_step_batched(
                ctx, y, kv_cache, sb.cross_mask_in,
                b.cross_q_w, b.cross_q_b, b.cross_out_w, b.cross_out_b,
                n_heads, d_model, i, T_enc_max, B);
            x = ggml_add(ctx, x, co);
        }
        {
            ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, b.norm_ffn_b);
            y = ffn(ctx, y, b.ffn_fc1_w, b.ffn_fc1_b, b.ffn_fc2_w, b.ffn_fc2_b);
            x = ggml_add(ctx, x, y);
        }
    }

    x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b);

    // Tied lm_head, no bias. Output: raw pre-softmax logits [vocab, B].
    ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x);
    ggml_set_name(logits, "step.logits_raw_b");
    ggml_set_output(logits);
    sb.logits_out = logits;
    ggml_build_forward_expand(sb.graph, sb.logits_out);
    return sb;
}

} // namespace transcribe::whisper