transcribe-cpp-sys 0.0.7

// arch/granite/weights.h - Granite Speech tensor catalog and hparams.
//
// INTERNAL to src/arch/granite/. Defines:
//
//   - GraniteHParams: architecture KV that drives tensor shapes. Read
//     from stt.granite.* and stt.frontend.* at load time, before any
//     tensors are allocated.
//
//   - GraniteWeights: named ggml_tensor* slots, one per logical
//     weight. Borrowed pointers — the tensors themselves live in the
//     model's ctx_meta / backend buffer.
//
// Architecture pattern: audio-llm. Three-component weight layout
// (mirrors the HF state-dict prefix scheme that the converter emits):
//
//   enc.*   Conformer encoder (16 layers, hidden=1024, Shaw block-local
//           attn over context_size=200 blocks, GLU conv with BatchNorm,
//           self-conditioned CTC bypass via out_mid at layer N/2)
//   proj.*  BLIP-2 Q-Former projector (3 learned queries, 2 layers of
//           self+cross attention + FFN, final LN, linear lift to
//           LM hidden)
//   dec.*   Granite-4 causal LM (40 layers, hidden=2048, GQA 16/4,
//           RMSNorm, SwiGLU, 4 scalar multipliers baked into the graph)
//
// Tensor naming mirrors the converter's output. Linear weights use
// PyTorch [out, in] order (i.e. ggml ne[0]=in, ne[1]=out). Conv1d
// kernels use PyTorch (out, in/groups, k) order (ggml ne[0]=k,
// ne[1]=in/groups, ne[2]=out).

#pragma once

#include "transcribe.h"

#include <cstdint>
#include <string>
#include <vector>

struct gguf_context;
struct ggml_context;
struct ggml_tensor;

namespace transcribe::granite {

// ---------------------------------------------------------------------------
// Hyperparameters
// ---------------------------------------------------------------------------

struct GraniteHParams {
    // Identity. `stt.variant` distinguishes models that share the same
    // architecture and hparams but differ in their published evaluation
    // prompt (e.g. granite-4.0-1b-speech, granite-speech-4.1-2b,
    // granite-speech-4.1-2b-plus). Used in model.cpp to pick the
    // per-variant instruction + optional system prompt that matches the
    // model card. Empty string falls back to the historical default.
    std::string variant;

    // Encoder (Conformer, granite_speech_encoder).
    int32_t enc_n_layers         = 0;
    int32_t enc_hidden           = 0;  // 1024
    int32_t enc_n_heads          = 0;  // 8
    int32_t enc_head_dim         = 0;  // 128
    int32_t enc_input_dim        = 0;  // 160 (= 2 mel frames stacked)
    int32_t enc_output_dim       = 0;  // 348 (CTC vocabulary head)
    int32_t enc_feedforward_mult = 0;  // 4 (FFN inner = hidden*mult)
    int32_t enc_conv_kernel_size = 0;  // 15
    int32_t enc_conv_expansion   = 0;  // 2 (conv inner = hidden*expansion)
    int32_t enc_max_pos_emb      = 0;  // 512
    int32_t enc_context_size     = 0;  // 200 (block size for local attn)

    // cat_hidden_layers (uint32 array, possibly empty). For -plus this
    // is [3]: the encoder concatenates block[3].out with the final
    // block hidden along the channel axis, doubling the projector's
    // cross-attention K/V input from 1024 to 2048. Empty list means
    // no concat (projector_input == hidden).
    std::vector<int32_t> enc_cat_hidden_layers;

    // Projector (BLIP-2 Q-Former).
    int32_t prj_n_layers            = 0;  // 2
    int32_t prj_hidden              = 0;  // 1024
    int32_t prj_intermediate        = 0;  // 4096
    int32_t prj_n_heads             = 0;  // 16
    int32_t prj_encoder_hidden_size = 0;  // 1024 or 2048 (varies on -plus)
    int32_t prj_cross_attn_freq     = 0;  // 1 (every layer has cross-attn)
    std::string prj_hidden_act;            // "gelu"
    float   prj_layer_norm_eps      = 0.0f;
    int32_t prj_max_pos_emb         = 0;  // 2048 (Q-Former absolute pos emb)
    std::string prj_pos_embed_type;        // "absolute"

    // Text LM (Granite-4 causal LM).
    int32_t dec_n_layers     = 0;  // 40
    int32_t dec_hidden       = 0;  // 2048
    int32_t dec_intermediate = 0;  // 4096
    int32_t dec_n_heads      = 0;  // 16
    int32_t dec_n_kv_heads   = 0;  // 4 (GQA 16/4)
    int32_t dec_head_dim     = 0;  // 128
    std::string dec_hidden_act;    // "silu"
    float   dec_rms_norm_eps = 0.0f;
    float   dec_rope_theta   = 0.0f;
    int32_t dec_max_position_embeddings = 0;  // 4096
    bool    dec_tie_word_embeddings     = false;  // true on -plus only
    int32_t dec_vocab_size              = 0;  // 100353

    // Granite-4 scalar multipliers. Each one is silently load-bearing —
    // missing any of them degrades accuracy without crashing.
    float dec_embedding_multiplier = 0.0f;  // 12.0
    float dec_logits_scaling       = 0.0f;  //  8.0  (logits /= this)
    float dec_attention_multiplier = 0.0f;  //  1/128 = 0.0078125 (softmax scale)
    float dec_residual_multiplier  = 0.0f;  //  0.22 (every residual add)

    // Audio fusion.
    int32_t audio_token_id  = 0;  // 100352
    int32_t downsample_rate = 0;  // 5
    int32_t window_size     = 0;  // 15 (encoder frames per Q-Former window)

    // Resolved token ids (filled from tokenizer KV at load).
    int32_t bos_token_id = -1;   // 100257
    int32_t eos_token_id = -1;   // 100257 (same as bos)
    int32_t pad_token_id = -1;   // 100256
    int32_t vocab_size   = 0;    // 100353

    // Frontend (torchaudio MelSpectrogram). htk-norm mel, no log, no
    // per-utterance norm, power=2 spectrogram. Stage 2b extends
    // MelFrontend to support these knobs; for now we just record them.
    std::string fe_type;             // "mel"
    int32_t     fe_num_mels      = 0;  // 80
    int32_t     fe_sample_rate   = 0;  // 16000
    int32_t     fe_n_fft         = 0;  // 512
    int32_t     fe_win_length    = 0;  // 400
    int32_t     fe_hop_length    = 0;  // 160
    std::string fe_window;             // "hann_periodic"
    std::string fe_normalize;          // "none"
    float       fe_dither        = 0.0f;
    float       fe_pre_emphasis  = 0.0f;
    float       fe_f_min         = 0.0f;
    float       fe_f_max         = 0.0f;
    std::string fe_pad_mode;           // "reflect"
    bool        fe_center        = true;
    std::string fe_mel_norm;           // "htk"
};

transcribe_status read_granite_hparams(const gguf_context * gguf,
                                       GraniteHParams &     hp);

// ---------------------------------------------------------------------------
// Weight slots - encoder (Conformer)
// ---------------------------------------------------------------------------

// Top-level encoder tensors: input linear (160 -> 1024) and the
// self-conditioned CTC bypass pair (ctc_proj: 1024 -> 348 for the
// auxiliary CTC head; ctc_bypass: 348 -> 1024 to add the softmaxed CTC
// logits back into the residual stream at layer N/2).
struct GraniteEncTop {
    ggml_tensor * input_linear_w = nullptr;  // [input_dim,  hidden]
    ggml_tensor * input_linear_b = nullptr;  // [hidden]
    ggml_tensor * ctc_proj_w     = nullptr;  // [hidden,     output_dim]
    ggml_tensor * ctc_proj_b     = nullptr;  // [output_dim]
    ggml_tensor * ctc_bypass_w   = nullptr;  // [output_dim, hidden]
    ggml_tensor * ctc_bypass_b   = nullptr;  // [hidden]
};

// One Conformer block. Macaron FFN (ff1, ff2) flank Shaw block-local
// self-attention and a GLU conv module with BatchNorm. Final per-block
// post-LN. attn.q is separate; attn.kv is fused (2*inner_dim output);
// attn.rel_pos_emb is a learned table for Shaw-style positional bias.
struct GraniteEncBlock {
    // FF1 (first half macaron). LN gamma/beta + up/down projections,
    // both with biases. SiLU between.
    ggml_tensor * norm_ff1_w   = nullptr;  // [hidden]
    ggml_tensor * norm_ff1_b   = nullptr;
    ggml_tensor * ff1_up_w     = nullptr;  // [hidden, hidden*ffn_mult]
    ggml_tensor * ff1_up_b     = nullptr;  // [hidden*ffn_mult]
    ggml_tensor * ff1_down_w   = nullptr;  // [hidden*ffn_mult, hidden]
    ggml_tensor * ff1_down_b   = nullptr;  // [hidden]

    // Block-local Shaw self-attention.
    ggml_tensor * norm_attn_w  = nullptr;  // [hidden]
    ggml_tensor * norm_attn_b  = nullptr;
    ggml_tensor * attn_q_w     = nullptr;  // [hidden, inner_dim]
    ggml_tensor * attn_kv_w    = nullptr;  // [hidden, 2*inner_dim]  (fused K|V)
    ggml_tensor * attn_out_w   = nullptr;  // [inner_dim, hidden]
    ggml_tensor * attn_out_b   = nullptr;  // [hidden]
    // Shaw relative-position table: [head_dim, 2*max_pos_emb+1].
    // Looked up by clamped (k - q + max_pos_emb) and added to QK^T.
    ggml_tensor * attn_rel_pos_emb = nullptr;  // [head_dim, 2*max_pos_emb+1]

    // Conv module: LN -> pointwise expand to 2*inner_dim -> GLU split
    // -> depthwise(k=15) -> BN -> SiLU -> pointwise contract back to
    // hidden. inner_dim = hidden * conv_expansion.
    ggml_tensor * norm_conv_w        = nullptr;  // [hidden]
    ggml_tensor * norm_conv_b        = nullptr;
    ggml_tensor * conv_pointwise1_w  = nullptr;  // [1, hidden, 2*inner_dim]
    ggml_tensor * conv_pointwise1_b  = nullptr;  // [2*inner_dim]
    ggml_tensor * conv_depthwise_w   = nullptr;  // [k, 1, inner_dim]
    ggml_tensor * conv_bn_w          = nullptr;  // [inner_dim]  (gamma)
    ggml_tensor * conv_bn_b          = nullptr;  // [inner_dim]  (beta)
    ggml_tensor * conv_bn_mean       = nullptr;  // [inner_dim]  (running mean)
    ggml_tensor * conv_bn_var        = nullptr;  // [inner_dim]  (running var)
    ggml_tensor * conv_pointwise2_w  = nullptr;  // [1, inner_dim, hidden]
    ggml_tensor * conv_pointwise2_b  = nullptr;  // [hidden]
    // Fused BatchNorm scale/bias precomputed at load (model.cpp
    // fuse_batch_norm). Replaces the raw weight/bias + running stats
    // for graph use; the raw tensors are still loaded but only feed
    // the fusion step. Both are [inner_dim] f32, allocated in a
    // separate ctx on the model.
    ggml_tensor * conv_bn_fused_scale = nullptr;  // [inner_dim]
    ggml_tensor * conv_bn_fused_bias  = nullptr;  // [inner_dim]

    // FF2 (second half macaron).
    ggml_tensor * norm_ff2_w   = nullptr;
    ggml_tensor * norm_ff2_b   = nullptr;
    ggml_tensor * ff2_up_w     = nullptr;
    ggml_tensor * ff2_up_b     = nullptr;
    ggml_tensor * ff2_down_w   = nullptr;
    ggml_tensor * ff2_down_b   = nullptr;

    // Post-block LayerNorm.
    ggml_tensor * norm_post_w  = nullptr;
    ggml_tensor * norm_post_b  = nullptr;
};

// ---------------------------------------------------------------------------
// Weight slots - projector (BLIP-2 Q-Former)
// ---------------------------------------------------------------------------

// Top-level projector tensors.
//   query: [hidden, num_queries=3, 1] — learned queries broadcast over
//          per-window encoder slices.
//   linear: lifts Q-Former hidden (1024) to LM hidden (2048).
//   qformer.final_norm: post Q-Former-stack LN over the query stream.
struct GraniteProjTop {
    ggml_tensor * query           = nullptr;  // [hidden, num_queries, 1]
    ggml_tensor * linear_w        = nullptr;  // [hidden, lm_hidden]
    ggml_tensor * linear_b        = nullptr;  // [lm_hidden]
    ggml_tensor * qformer_final_norm_w = nullptr;  // [hidden]
    ggml_tensor * qformer_final_norm_b = nullptr;
};

// One Q-Former layer: self-attn -> LN -> cross-attn -> LN -> FFN -> LN.
// All linears carry bias. The cross-attention K/V input dim is
// `encoder_hidden_size`, which is 1024 for the 1b/2b variants and 2048
// for the -plus variant.
struct GraniteProjBlock {
    // Self-attention (Q/K/V all hidden -> hidden).
    ggml_tensor * self_attn_q_w   = nullptr;
    ggml_tensor * self_attn_q_b   = nullptr;
    ggml_tensor * self_attn_k_w   = nullptr;
    ggml_tensor * self_attn_k_b   = nullptr;
    ggml_tensor * self_attn_v_w   = nullptr;
    ggml_tensor * self_attn_v_b   = nullptr;
    ggml_tensor * self_attn_out_w = nullptr;
    ggml_tensor * self_attn_out_b = nullptr;
    ggml_tensor * norm_self_attn_w = nullptr;
    ggml_tensor * norm_self_attn_b = nullptr;

    // Cross-attention. Q from query stream; K/V from encoder slice.
    ggml_tensor * cross_attn_q_w   = nullptr;  // [prj_hidden,        prj_hidden]
    ggml_tensor * cross_attn_q_b   = nullptr;
    ggml_tensor * cross_attn_k_w   = nullptr;  // [encoder_hidden,    prj_hidden]
    ggml_tensor * cross_attn_k_b   = nullptr;
    ggml_tensor * cross_attn_v_w   = nullptr;  // [encoder_hidden,    prj_hidden]
    ggml_tensor * cross_attn_v_b   = nullptr;
    ggml_tensor * cross_attn_out_w = nullptr;  // [prj_hidden,        prj_hidden]
    ggml_tensor * cross_attn_out_b = nullptr;
    ggml_tensor * norm_cross_attn_w = nullptr;
    ggml_tensor * norm_cross_attn_b = nullptr;

    // FFN (intermediate_query / output_query collapse).
    ggml_tensor * ffn_up_w    = nullptr;  // [hidden, intermediate]
    ggml_tensor * ffn_up_b    = nullptr;
    ggml_tensor * ffn_down_w  = nullptr;  // [intermediate, hidden]
    ggml_tensor * ffn_down_b  = nullptr;
    ggml_tensor * norm_ffn_w  = nullptr;
    ggml_tensor * norm_ffn_b  = nullptr;
};

// ---------------------------------------------------------------------------
// Weight slots - Granite-4 LM
// ---------------------------------------------------------------------------

struct GraniteDecEmbed {
    ggml_tensor * token_w = nullptr;  // [hidden, vocab]
};

struct GraniteDecBlock {
    ggml_tensor * norm_attn_w = nullptr;  // input_layernorm (RMSNorm, no bias)
    ggml_tensor * norm_ffn_w  = nullptr;  // post_attention_layernorm
    // GQA projections; no biases on Granite-4.
    ggml_tensor * attn_q_w    = nullptr;  // [hidden, n_heads * head_dim]
    ggml_tensor * attn_k_w    = nullptr;  // [hidden, n_kv_heads * head_dim]
    ggml_tensor * attn_v_w    = nullptr;  // [hidden, n_kv_heads * head_dim]
    ggml_tensor * attn_o_w    = nullptr;  // [n_heads * head_dim, hidden]
    // SwiGLU MLP.
    ggml_tensor * ffn_gate_w  = nullptr;  // [hidden, intermediate]
    ggml_tensor * ffn_up_w    = nullptr;  // [hidden, intermediate]
    ggml_tensor * ffn_down_w  = nullptr;  // [intermediate, hidden]
    // Optional gate/up pack: [hidden, 2*intermediate]. Allocated by
    // causal_lm::pack_gate_up at load time.
    ggml_tensor * ffn_gate_up_w = nullptr;
};

struct GraniteDecFinal {
    ggml_tensor * norm_w   = nullptr;   // RMSNorm before lm_head
    // lm_head, only present when tie_word_embeddings=false (1b / 2b
    // variants). When tied (plus), the graph reuses dec_embed.token_w.
    ggml_tensor * output_w = nullptr;   // [hidden, vocab] or nullptr
};

// ---------------------------------------------------------------------------
// Aggregated weights
// ---------------------------------------------------------------------------

struct GraniteWeights {
    GraniteEncTop                   enc_top;
    std::vector<GraniteEncBlock>    enc_blocks;

    GraniteProjTop                  proj_top;
    std::vector<GraniteProjBlock>   proj_blocks;

    GraniteDecEmbed                 dec_embed;
    std::vector<GraniteDecBlock>    dec_blocks;
    GraniteDecFinal                 dec_final;
};

transcribe_status build_granite_weights(ggml_context *         ctx_meta,
                                        const GraniteHParams & hp,
                                        GraniteWeights &       weights);

} // namespace transcribe::granite