transcribe-cpp-sys 0.0.7

// arch/granite/weights.cpp - read_granite_hparams + build_granite_weights.
//
// Pattern mirrors arch/qwen3_asr/weights.cpp. Every required KV is
// read explicitly; BadType is always fatal. The tensor catalog is a
// sequence of find_tensor() calls with expected shapes; missing tensor
// or shape mismatch returns TRANSCRIBE_ERR_GGUF.

#include "weights.h"

#include "transcribe-meta.h"
#include "transcribe-weights-util.h"
#include "transcribe-log.h"

#include "ggml.h"
#include "gguf.h"

#include <cstdio>

namespace transcribe::granite {

namespace {
constexpr const char * kFamilyTag = "granite";
}

transcribe_status read_granite_hparams(const gguf_context * gguf,
                                       GraniteHParams &     hp)
{
    if (gguf == nullptr) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }

    // ----- Identity -----
    if (auto st = read_optional_string_kv(
            gguf, "stt.variant", kFamilyTag, "", hp.variant);
        st != TRANSCRIBE_OK) return st;

    // ----- Encoder (Conformer) -----
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.n_layers", kFamilyTag, hp.enc_n_layers); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.hidden", kFamilyTag, hp.enc_hidden); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.n_heads", kFamilyTag, hp.enc_n_heads); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.head_dim", kFamilyTag, hp.enc_head_dim); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.input_dim", kFamilyTag, hp.enc_input_dim); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.output_dim", kFamilyTag, hp.enc_output_dim); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.feedforward_mult", kFamilyTag, hp.enc_feedforward_mult); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.conv_kernel_size", kFamilyTag, hp.enc_conv_kernel_size); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.conv_expansion", kFamilyTag, hp.enc_conv_expansion); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.max_pos_emb", kFamilyTag, hp.enc_max_pos_emb); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.context_size", kFamilyTag, hp.enc_context_size); st != TRANSCRIBE_OK) return st;

    // cat_hidden_layers is an optional array. Absent or empty array both
    // mean "no concat". Converter writes the uint32 array via gguf-py;
    // we read via read_int32_array_kv (uint32 ↔ int32 is safe for layer
    // indices in [0, num_layers)).
    {
        std::vector<int32_t> tmp;
        switch (read_int32_array_kv(gguf, "stt.granite.encoder.cat_hidden_layers", tmp)) {
            case KvResult::Ok:      hp.enc_cat_hidden_layers = std::move(tmp); break;
            case KvResult::Absent:  hp.enc_cat_hidden_layers.clear();          break;
            case KvResult::BadType:
                log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                             "granite: stt.granite.encoder.cat_hidden_layers has wrong type");
                return TRANSCRIBE_ERR_GGUF;
        }
    }

    // ----- Projector (BLIP-2 Q-Former) -----
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.n_layers", kFamilyTag, hp.prj_n_layers); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.hidden", kFamilyTag, hp.prj_hidden); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.intermediate", kFamilyTag, hp.prj_intermediate); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.n_heads", kFamilyTag, hp.prj_n_heads); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.encoder_hidden_size", kFamilyTag, hp.prj_encoder_hidden_size); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.cross_attn_frequency", kFamilyTag, hp.prj_cross_attn_freq); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.granite.projector.hidden_act", kFamilyTag, hp.prj_hidden_act); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.projector.layer_norm_eps", kFamilyTag, hp.prj_layer_norm_eps); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.max_pos_emb", kFamilyTag, hp.prj_max_pos_emb); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.granite.projector.position_embedding_type", kFamilyTag, hp.prj_pos_embed_type); st != TRANSCRIBE_OK) return st;

    // ----- Text LM (Granite-4) -----
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.n_layers", kFamilyTag, hp.dec_n_layers); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.hidden_size", kFamilyTag, hp.dec_hidden); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.intermediate_size", kFamilyTag, hp.dec_intermediate); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.n_heads", kFamilyTag, hp.dec_n_heads); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.n_kv_heads", kFamilyTag, hp.dec_n_kv_heads); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.head_dim", kFamilyTag, hp.dec_head_dim); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.granite.decoder.hidden_act", kFamilyTag, hp.dec_hidden_act); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.rms_norm_eps", kFamilyTag, hp.dec_rms_norm_eps); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.rope_theta", kFamilyTag, hp.dec_rope_theta); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.max_position_embeddings", kFamilyTag, hp.dec_max_position_embeddings); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv(gguf, "stt.granite.decoder.tie_word_embeddings", kFamilyTag, false, hp.dec_tie_word_embeddings); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.vocab_size", kFamilyTag, hp.dec_vocab_size); st != TRANSCRIBE_OK) return st;

    // Granite-4 scalar multipliers. All four are required — missing any
    // of them silently degrades accuracy without crashing.
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.embedding_multiplier", kFamilyTag, hp.dec_embedding_multiplier); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.logits_scaling", kFamilyTag, hp.dec_logits_scaling); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.attention_multiplier", kFamilyTag, hp.dec_attention_multiplier); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.residual_multiplier", kFamilyTag, hp.dec_residual_multiplier); st != TRANSCRIBE_OK) return st;

    // ----- Audio fusion -----
    if (auto st = read_required_u32_kv(gguf, "stt.granite.audio_token_id", kFamilyTag, hp.audio_token_id); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.downsample_rate", kFamilyTag, hp.downsample_rate); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.granite.window_size", kFamilyTag, hp.window_size); st != TRANSCRIBE_OK) return st;

    // ----- Frontend (torchaudio MelSpectrogram) -----
    if (auto st = read_required_string_kv(gguf, "stt.frontend.type", kFamilyTag, hp.fe_type); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.frontend.num_mels", kFamilyTag, hp.fe_num_mels); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.frontend.sample_rate", kFamilyTag, hp.fe_sample_rate); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.frontend.n_fft", kFamilyTag, hp.fe_n_fft); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.frontend.win_length", kFamilyTag, hp.fe_win_length); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_u32_kv(gguf, "stt.frontend.hop_length", kFamilyTag, hp.fe_hop_length); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.window", kFamilyTag, hp.fe_window); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_string_kv(gguf, "stt.frontend.normalize", kFamilyTag, hp.fe_normalize); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.frontend.dither", kFamilyTag, hp.fe_dither); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.frontend.pre_emphasis", kFamilyTag, hp.fe_pre_emphasis); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.frontend.f_min", kFamilyTag, hp.fe_f_min); st != TRANSCRIBE_OK) return st;
    if (auto st = read_required_f32_kv(gguf, "stt.frontend.f_max", kFamilyTag, hp.fe_f_max); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_string_kv(gguf, "stt.frontend.pad_mode", kFamilyTag, "reflect", hp.fe_pad_mode); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_string_kv(gguf, "stt.frontend.mel_norm", kFamilyTag, "htk", hp.fe_mel_norm); st != TRANSCRIBE_OK) return st;
    if (auto st = read_optional_bool_kv(gguf, "stt.frontend.center", kFamilyTag, true, hp.fe_center); st != TRANSCRIBE_OK) return st;

    // ----- Cross-field invariants -----
    if (hp.enc_n_layers <= 0 || hp.enc_hidden <= 0 || hp.enc_n_heads <= 0 ||
        hp.enc_head_dim <= 0 || hp.enc_input_dim <= 0 || hp.enc_output_dim <= 0 ||
        hp.enc_feedforward_mult <= 0 || hp.enc_conv_kernel_size <= 0 ||
        hp.enc_conv_expansion <= 0 || hp.enc_max_pos_emb <= 0 ||
        hp.enc_context_size <= 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite: encoder hparams must be positive");
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_n_heads * hp.enc_head_dim != hp.enc_hidden) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: encoder n_heads * head_dim (%d * %d) "
                     "!= hidden (%d)",
                     hp.enc_n_heads, hp.enc_head_dim, hp.enc_hidden);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.enc_context_size > hp.enc_max_pos_emb) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: context_size (%d) > max_pos_emb (%d)",
                     hp.enc_context_size, hp.enc_max_pos_emb);
        return TRANSCRIBE_ERR_GGUF;
    }
    // cat_hidden_layers entries must reference real layers and the
    // expanded projector input dim must match prj_encoder_hidden_size.
    {
        const int32_t expected_kv_dim = hp.enc_hidden *
            (static_cast<int32_t>(hp.enc_cat_hidden_layers.size()) + 1);
        if (expected_kv_dim != hp.prj_encoder_hidden_size) {
            log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                         "granite: prj.encoder_hidden_size (%d) does not match "
                         "hidden*(1+|cat_hidden_layers|) = %d * (1 + %d) = %d",
                         hp.prj_encoder_hidden_size, hp.enc_hidden,
                         static_cast<int>(hp.enc_cat_hidden_layers.size()),
                         expected_kv_dim);
            return TRANSCRIBE_ERR_GGUF;
        }
        for (int32_t idx : hp.enc_cat_hidden_layers) {
            if (idx < 0 || idx >= hp.enc_n_layers) {
                log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                             "granite: cat_hidden_layers entry %d out of range "
                             "[0, %d)", idx, hp.enc_n_layers);
                return TRANSCRIBE_ERR_GGUF;
            }
        }
    }
    if (hp.prj_n_layers <= 0 || hp.prj_hidden <= 0 ||
        hp.prj_intermediate <= 0 || hp.prj_n_heads <= 0 ||
        hp.prj_encoder_hidden_size <= 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite: projector hparams must be positive");
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.prj_hidden % hp.prj_n_heads != 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: projector hidden (%d) not divisible by n_heads (%d)",
                     hp.prj_hidden, hp.prj_n_heads);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.prj_hidden_act != "gelu") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: unsupported projector hidden_act \"%s\" "
                     "(only gelu)", hp.prj_hidden_act.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.prj_pos_embed_type != "absolute") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: unsupported projector position_embedding_type "
                     "\"%s\" (only absolute)", hp.prj_pos_embed_type.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.dec_n_layers <= 0 || hp.dec_hidden <= 0 || hp.dec_n_heads <= 0 ||
        hp.dec_n_kv_heads <= 0 || hp.dec_head_dim <= 0 || hp.dec_intermediate <= 0)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite: decoder hparams must be positive");
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.dec_n_heads % hp.dec_n_kv_heads != 0) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: dec n_heads (%d) not divisible by n_kv_heads (%d)",
                     hp.dec_n_heads, hp.dec_n_kv_heads);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.dec_n_heads * hp.dec_head_dim != hp.dec_hidden) {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: dec n_heads * head_dim (%d * %d) != hidden (%d)",
                     hp.dec_n_heads, hp.dec_head_dim, hp.dec_hidden);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.dec_hidden_act != "silu" && hp.dec_hidden_act != "swish") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: unsupported decoder hidden_act \"%s\" "
                     "(only silu/swish)", hp.dec_hidden_act.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.dec_embedding_multiplier <= 0.0f ||
        hp.dec_logits_scaling       <= 0.0f ||
        hp.dec_attention_multiplier <= 0.0f ||
        hp.dec_residual_multiplier  <= 0.0f)
    {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: decoder multipliers must be > 0 "
                     "(emb=%g, logits=%g, attn=%g, residual=%g)",
                     hp.dec_embedding_multiplier, hp.dec_logits_scaling,
                     hp.dec_attention_multiplier, hp.dec_residual_multiplier);
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_type != "mel") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite: unsupported frontend type \"%s\"",
                     hp.fe_type.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_normalize != "per_utterance") {
        // GraniteSpeechFeatureExtractor applies whisper-style per-utterance
        // normalization (log10 → max-8 floor → /4 + 1). The C++ MelFrontend
        // produces a bit-identical result under "per_utterance" mode.
        // Other modes would mean the converter changed; fail loudly.
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: unsupported frontend normalize \"%s\" "
                     "(only \"per_utterance\")", hp.fe_normalize.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    if (hp.fe_mel_norm != "htk") {
        log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                     "granite: unsupported frontend mel_norm \"%s\" "
                     "(only \"htk\")", hp.fe_mel_norm.c_str());
        return TRANSCRIBE_ERR_GGUF;
    }
    // Audio injection precondition: projector lifts to LM hidden, so
    // the linear-lift output dim must equal dec.hidden.
    // (We check that on the actual tensor shape during build_weights.)

    return TRANSCRIBE_OK;
}

// ---------------------------------------------------------------------------
// Weights
// ---------------------------------------------------------------------------

namespace {

using transcribe::weights::lname;
constexpr const char * kTag = kFamilyTag;

#define GET_F32(slot, name, ...) \
    do { \
        ggml_tensor * _t = transcribe::weights::find_tensor( \
            ctx_meta, (name), \
            {GGML_TYPE_F32}, {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

#define GET_CONV(slot, name, ...) \
    do { \
        ggml_tensor * _t = transcribe::weights::find_tensor( \
            ctx_meta, (name), \
            {TRANSCRIBE_QUANT_CONV_TYPES}, \
            {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

#define GET_LIN(slot, name, ...) \
    do { \
        ggml_tensor * _t = transcribe::weights::find_tensor( \
            ctx_meta, (name), \
            {TRANSCRIBE_QUANT_LINEAR_TYPES}, \
            {__VA_ARGS__}, kTag); \
        if (_t == nullptr) return TRANSCRIBE_ERR_GGUF; \
        (slot) = _t; \
    } while (0)

} // namespace

transcribe_status build_granite_weights(ggml_context *         ctx_meta,
                                        const GraniteHParams & hp,
                                        GraniteWeights &       weights)
{
    if (ctx_meta == nullptr) {
        return TRANSCRIBE_ERR_INVALID_ARG;
    }

    // ----- Derived dims -----
    const int64_t enc_h        = hp.enc_hidden;
    const int64_t enc_in       = hp.enc_input_dim;
    const int64_t enc_out      = hp.enc_output_dim;
    const int64_t enc_n_heads  = hp.enc_n_heads;
    const int64_t enc_head_dim = hp.enc_head_dim;
    const int64_t enc_inner    = enc_n_heads * enc_head_dim;   // typically == enc_h
    const int64_t enc_ffn      = enc_h * hp.enc_feedforward_mult;
    const int64_t conv_inner   = enc_h * hp.enc_conv_expansion;
    const int64_t conv_up_out  = conv_inner * 2;               // pre-GLU
    const int64_t conv_k       = hp.enc_conv_kernel_size;
    const int64_t rel_pos_len  = 2 * hp.enc_max_pos_emb + 1;

    const int64_t prj_h        = hp.prj_hidden;
    const int64_t prj_im       = hp.prj_intermediate;
    const int64_t prj_kv_in    = hp.prj_encoder_hidden_size;   // 1024 or 2048
    const int64_t prj_n_layers = hp.prj_n_layers;
    // num_queries (3 on granite) is not in KV — it's the ne[1] of the
    // query tensor. Read it implicitly via find_tensor's ne[] match.

    const int64_t dec_h     = hp.dec_hidden;
    const int64_t dec_vocab = hp.dec_vocab_size;
    const int64_t dec_nh    = hp.dec_n_heads;
    const int64_t dec_nkv   = hp.dec_n_kv_heads;
    const int64_t dec_hd    = hp.dec_head_dim;
    const int64_t dec_im    = hp.dec_intermediate;
    const int64_t q_out     = dec_nh  * dec_hd;
    const int64_t kv_out    = dec_nkv * dec_hd;

    // ----- Encoder: top-level -----
    GET_LIN(weights.enc_top.input_linear_w, "enc.input_linear.weight", enc_in,  enc_h);
    GET_F32(weights.enc_top.input_linear_b, "enc.input_linear.bias",   enc_h);
    GET_LIN(weights.enc_top.ctc_proj_w,     "enc.ctc_proj.weight",     enc_h,   enc_out);
    GET_F32(weights.enc_top.ctc_proj_b,     "enc.ctc_proj.bias",       enc_out);
    GET_LIN(weights.enc_top.ctc_bypass_w,   "enc.ctc_bypass.weight",   enc_out, enc_h);
    GET_F32(weights.enc_top.ctc_bypass_b,   "enc.ctc_bypass.bias",     enc_h);

    // ----- Encoder: per-block -----
    weights.enc_blocks.assign(hp.enc_n_layers, GraniteEncBlock{});
    for (int i = 0; i < hp.enc_n_layers; ++i) {
        auto & b = weights.enc_blocks[i];

        // FF1 (macaron half).
        GET_F32(b.norm_ff1_w, lname("enc.blocks.%d.norm_ff1.weight", i), enc_h);
        GET_F32(b.norm_ff1_b, lname("enc.blocks.%d.norm_ff1.bias",   i), enc_h);
        GET_LIN(b.ff1_up_w,   lname("enc.blocks.%d.ff1.up.weight",   i), enc_h,   enc_ffn);
        GET_F32(b.ff1_up_b,   lname("enc.blocks.%d.ff1.up.bias",     i), enc_ffn);
        GET_LIN(b.ff1_down_w, lname("enc.blocks.%d.ff1.down.weight", i), enc_ffn, enc_h);
        GET_F32(b.ff1_down_b, lname("enc.blocks.%d.ff1.down.bias",   i), enc_h);

        // Block-local Shaw self-attention.
        GET_F32(b.norm_attn_w, lname("enc.blocks.%d.norm_attn.weight", i), enc_h);
        GET_F32(b.norm_attn_b, lname("enc.blocks.%d.norm_attn.bias",   i), enc_h);
        // attn.q : [enc_h, inner_dim]. No upstream bias.
        GET_LIN(b.attn_q_w,    lname("enc.blocks.%d.attn.q.weight",    i), enc_h, enc_inner);
        // attn.kv fused output dim is 2*inner_dim. No bias.
        GET_LIN(b.attn_kv_w,   lname("enc.blocks.%d.attn.kv.weight",   i), enc_h, 2 * enc_inner);
        // attn.out has bias.
        GET_LIN(b.attn_out_w,  lname("enc.blocks.%d.attn.out.weight",  i), enc_inner, enc_h);
        GET_F32(b.attn_out_b,  lname("enc.blocks.%d.attn.out.bias",    i), enc_h);
        // Shaw rel-pos embedding table. find_tensor accepts any linear
        // dtype here (BF16/F16/F32) — the loader will keep whatever the
        // converter emitted.
        GET_LIN(b.attn_rel_pos_emb, lname("enc.blocks.%d.attn.rel_pos_emb.weight", i),
                enc_head_dim, rel_pos_len);

        // Conv module (LN -> pointwise expand -> GLU -> depthwise -> BN -> SiLU -> pointwise contract).
        GET_F32(b.norm_conv_w, lname("enc.blocks.%d.norm_conv.weight", i), enc_h);
        GET_F32(b.norm_conv_b, lname("enc.blocks.%d.norm_conv.bias",   i), enc_h);
        // Conv1d kernel layout (k, in/groups, out_channels). pointwise k=1.
        GET_CONV(b.conv_pointwise1_w, lname("enc.blocks.%d.conv.pointwise1.weight", i),
                 1, enc_h, conv_up_out);
        GET_F32(b.conv_pointwise1_b,  lname("enc.blocks.%d.conv.pointwise1.bias",   i),
                conv_up_out);
        // Depthwise: groups = in_channels = conv_inner.
        GET_CONV(b.conv_depthwise_w,  lname("enc.blocks.%d.conv.depthwise.weight",  i),
                 conv_k, 1, conv_inner);
        // BatchNorm1d fields are F32 (the converter routes anything
        // matching .bn. into the F32 bucket).
        GET_F32(b.conv_bn_w,    lname("enc.blocks.%d.conv.bn.weight",       i), conv_inner);
        GET_F32(b.conv_bn_b,    lname("enc.blocks.%d.conv.bn.bias",         i), conv_inner);
        GET_F32(b.conv_bn_mean, lname("enc.blocks.%d.conv.bn.running_mean", i), conv_inner);
        GET_F32(b.conv_bn_var,  lname("enc.blocks.%d.conv.bn.running_var",  i), conv_inner);
        GET_CONV(b.conv_pointwise2_w, lname("enc.blocks.%d.conv.pointwise2.weight", i),
                 1, conv_inner, enc_h);
        GET_F32(b.conv_pointwise2_b,  lname("enc.blocks.%d.conv.pointwise2.bias",   i),
                enc_h);

        // FF2 (macaron half).
        GET_F32(b.norm_ff2_w, lname("enc.blocks.%d.norm_ff2.weight", i), enc_h);
        GET_F32(b.norm_ff2_b, lname("enc.blocks.%d.norm_ff2.bias",   i), enc_h);
        GET_LIN(b.ff2_up_w,   lname("enc.blocks.%d.ff2.up.weight",   i), enc_h,   enc_ffn);
        GET_F32(b.ff2_up_b,   lname("enc.blocks.%d.ff2.up.bias",     i), enc_ffn);
        GET_LIN(b.ff2_down_w, lname("enc.blocks.%d.ff2.down.weight", i), enc_ffn, enc_h);
        GET_F32(b.ff2_down_b, lname("enc.blocks.%d.ff2.down.bias",   i), enc_h);

        // Post-block LN.
        GET_F32(b.norm_post_w, lname("enc.blocks.%d.norm_post.weight", i), enc_h);
        GET_F32(b.norm_post_b, lname("enc.blocks.%d.norm_post.bias",   i), enc_h);
    }

    // ----- Projector: top-level -----
    //
    // proj.query has shape (1, num_queries, hidden) in PyTorch, which
    // is ne[hidden, num_queries, 1] in ggml. num_queries is not in KV;
    // we accept whatever the file declared and stash it for the graph
    // builder via the tensor's ne[1]. The converter ships 3 queries on
    // every public variant. Use find_tensor with a wildcard num_queries
    // by not specifying it — but find_tensor requires every dim, so we
    // hardcode 3 here and add an explicit check below if a future
    // variant changes the count. The query tensor is a learned linear
    // embedding; we accept LINEAR_TYPES dtypes.
    GET_LIN(weights.proj_top.query,    "proj.query",    prj_h, 3, 1);
    GET_LIN(weights.proj_top.linear_w, "proj.linear.weight", prj_h, dec_h);
    GET_F32(weights.proj_top.linear_b, "proj.linear.bias",   dec_h);
    GET_F32(weights.proj_top.qformer_final_norm_w,
            "proj.qformer.final_norm.weight", prj_h);
    GET_F32(weights.proj_top.qformer_final_norm_b,
            "proj.qformer.final_norm.bias",   prj_h);

    // ----- Projector: per-Q-Former-layer -----
    weights.proj_blocks.assign(prj_n_layers, GraniteProjBlock{});
    for (int i = 0; i < prj_n_layers; ++i) {
        auto & b = weights.proj_blocks[i];

        // Self-attention: Q/K/V/Out all map prj_h -> prj_h, every linear
        // has bias.
        GET_LIN(b.self_attn_q_w, lname("proj.qformer.blocks.%d.self_attn.q.weight", i), prj_h, prj_h);
        GET_F32(b.self_attn_q_b, lname("proj.qformer.blocks.%d.self_attn.q.bias",   i), prj_h);
        GET_LIN(b.self_attn_k_w, lname("proj.qformer.blocks.%d.self_attn.k.weight", i), prj_h, prj_h);
        GET_F32(b.self_attn_k_b, lname("proj.qformer.blocks.%d.self_attn.k.bias",   i), prj_h);
        GET_LIN(b.self_attn_v_w, lname("proj.qformer.blocks.%d.self_attn.v.weight", i), prj_h, prj_h);
        GET_F32(b.self_attn_v_b, lname("proj.qformer.blocks.%d.self_attn.v.bias",   i), prj_h);
        GET_LIN(b.self_attn_out_w, lname("proj.qformer.blocks.%d.self_attn.out.weight", i), prj_h, prj_h);
        GET_F32(b.self_attn_out_b, lname("proj.qformer.blocks.%d.self_attn.out.bias",   i), prj_h);
        GET_F32(b.norm_self_attn_w, lname("proj.qformer.blocks.%d.norm_self_attn.weight", i), prj_h);
        GET_F32(b.norm_self_attn_b, lname("proj.qformer.blocks.%d.norm_self_attn.bias",   i), prj_h);

        // Cross-attention: Q is prj_h->prj_h; K/V are prj_kv_in->prj_h
        // (prj_kv_in == encoder hidden for 1b/2b, doubled for -plus).
        GET_LIN(b.cross_attn_q_w, lname("proj.qformer.blocks.%d.cross_attn.q.weight", i), prj_h,     prj_h);
        GET_F32(b.cross_attn_q_b, lname("proj.qformer.blocks.%d.cross_attn.q.bias",   i), prj_h);
        GET_LIN(b.cross_attn_k_w, lname("proj.qformer.blocks.%d.cross_attn.k.weight", i), prj_kv_in, prj_h);
        GET_F32(b.cross_attn_k_b, lname("proj.qformer.blocks.%d.cross_attn.k.bias",   i), prj_h);
        GET_LIN(b.cross_attn_v_w, lname("proj.qformer.blocks.%d.cross_attn.v.weight", i), prj_kv_in, prj_h);
        GET_F32(b.cross_attn_v_b, lname("proj.qformer.blocks.%d.cross_attn.v.bias",   i), prj_h);
        GET_LIN(b.cross_attn_out_w, lname("proj.qformer.blocks.%d.cross_attn.out.weight", i), prj_h, prj_h);
        GET_F32(b.cross_attn_out_b, lname("proj.qformer.blocks.%d.cross_attn.out.bias",   i), prj_h);
        GET_F32(b.norm_cross_attn_w, lname("proj.qformer.blocks.%d.norm_cross_attn.weight", i), prj_h);
        GET_F32(b.norm_cross_attn_b, lname("proj.qformer.blocks.%d.norm_cross_attn.bias",   i), prj_h);

        // FFN.
        GET_LIN(b.ffn_up_w,   lname("proj.qformer.blocks.%d.ffn.up.weight",   i), prj_h,  prj_im);
        GET_F32(b.ffn_up_b,   lname("proj.qformer.blocks.%d.ffn.up.bias",     i), prj_im);
        GET_LIN(b.ffn_down_w, lname("proj.qformer.blocks.%d.ffn.down.weight", i), prj_im, prj_h);
        GET_F32(b.ffn_down_b, lname("proj.qformer.blocks.%d.ffn.down.bias",   i), prj_h);
        GET_F32(b.norm_ffn_w, lname("proj.qformer.blocks.%d.norm_ffn.weight", i), prj_h);
        GET_F32(b.norm_ffn_b, lname("proj.qformer.blocks.%d.norm_ffn.bias",   i), prj_h);
    }

    // ----- Text LM: embed -----
    GET_LIN(weights.dec_embed.token_w, "dec.token_embd.weight", dec_h, dec_vocab);

    // ----- Text LM: blocks -----
    weights.dec_blocks.assign(hp.dec_n_layers, GraniteDecBlock{});
    for (int i = 0; i < hp.dec_n_layers; ++i) {
        auto & b = weights.dec_blocks[i];
        GET_F32(b.norm_attn_w, lname("dec.blocks.%d.norm_attn.weight", i), dec_h);
        GET_F32(b.norm_ffn_w,  lname("dec.blocks.%d.norm_ffn.weight",  i), dec_h);
        GET_LIN(b.attn_q_w,    lname("dec.blocks.%d.attn.q.weight",    i), dec_h, q_out);
        GET_LIN(b.attn_k_w,    lname("dec.blocks.%d.attn.k.weight",    i), dec_h, kv_out);
        GET_LIN(b.attn_v_w,    lname("dec.blocks.%d.attn.v.weight",    i), dec_h, kv_out);
        GET_LIN(b.attn_o_w,    lname("dec.blocks.%d.attn.o.weight",    i), q_out, dec_h);
        GET_LIN(b.ffn_gate_w,  lname("dec.blocks.%d.ffn.gate.weight",  i), dec_h, dec_im);
        GET_LIN(b.ffn_up_w,    lname("dec.blocks.%d.ffn.up.weight",    i), dec_h, dec_im);
        GET_LIN(b.ffn_down_w,  lname("dec.blocks.%d.ffn.down.weight",  i), dec_im, dec_h);

        if (b.ffn_gate_w->type != b.ffn_up_w->type) {
            log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                         "granite: ffn gate/up dtype mismatch at dec layer %d", i);
            return TRANSCRIBE_ERR_GGUF;
        }
    }

    // ----- Text LM: final norm + (conditional) output -----
    GET_F32(weights.dec_final.norm_w, "dec.output_norm.weight", dec_h);
    if (hp.dec_tie_word_embeddings) {
        // Tied head (granite-speech-4.1-2b-plus): no dec.output.weight
        // in the GGUF. The graph reuses dec_embed.token_w for the head.
        weights.dec_final.output_w = nullptr;
    } else {
        GET_LIN(weights.dec_final.output_w, "dec.output.weight", dec_h, dec_vocab);
    }

    return TRANSCRIBE_OK;
}

} // namespace transcribe::granite