transcribe-cpp-sys 0.2.4

// src/conformer/conformer.h - shared Conformer encoder helpers.
//
// Family-agnostic helpers for the NeMo-style Conformer encoder shared by
// Parakeet, Cohere-ASR, and future variants. Block topology:
//
//   pre_encode (DwStridingSubsampling) ->
//   N x { macaron FF1 (0.5x) -> rel-pos MHSA -> conv module
//         (pw-GLU -> depthwise -> BN -> SiLU -> pw) -> macaron FF2 (0.5x)
//         -> per-block LayerNorm }
//
// Driven by three structs: PreEncodeView / BlockView (nullable-pointer
// projections; null bias = skip the add), ConvPolicy (per-family pointwise/
// depthwise conv dispatch), and BlockParams (scalar knobs). Sub-helpers are
// exposed at namespace scope for families that hand-build a block for
// debug-dump. Helpers name no intermediates except build_pre_encode (explicit
// prefix). The f32-friendly conv helpers (conv_1d_f32, conv_2d_dw_f32,
// conv_1d_dw_f32) work around ggml kernel limitations on Metal (see
// conv_2d_dw_f32 in conformer.cpp); use them for f32 kernels on Metal.

#pragma once

#include "ggml.h"

namespace transcribe::conformer {

// LayerNorm + BatchNorm epsilon. Both default to 1e-5 in NeMo and are NOT
// stored in the GGUF, so hardcode to match.
constexpr float kLayerNormEps = 1e-5f;

// Nullable-pointer projection over the pre_encode weights. The family
// fills this from its own weights struct at the call site.
struct PreEncodeView {
    ggml_tensor * conv0_w = nullptr;  // [KW=3, KH=3, IC=1,        OC=channels]
    ggml_tensor * conv0_b = nullptr;  // [channels]
    ggml_tensor * conv2_w = nullptr;  // depthwise [KW=3, KH=3, 1, OC=channels]
    ggml_tensor * conv2_b = nullptr;  // [channels]
    ggml_tensor * conv3_w = nullptr;  // pointwise [1, 1, channels, channels]
    ggml_tensor * conv3_b = nullptr;  // [channels]
    ggml_tensor * conv5_w = nullptr;  // depthwise
    ggml_tensor * conv5_b = nullptr;
    ggml_tensor * conv6_w = nullptr;  // pointwise
    ggml_tensor * conv6_b = nullptr;
    // Final projection from [channels * (num_mels / subsampling)] to d_model.
    ggml_tensor * out_w   = nullptr;  // [pre_encode_in, d_model]
    ggml_tensor * out_b   = nullptr;  // [d_model]
};

// Nullable-pointer projection over one Conformer block's weights.
// Every `_b` (bias) slot is optional — null = skip the add.
struct BlockView {
    // Macaron FF1.
    ggml_tensor * norm_ff1_w = nullptr;  // [d_model]
    ggml_tensor * norm_ff1_b = nullptr;  // [d_model]
    ggml_tensor * ff1_lin1_w = nullptr;  // [d_model, d_ff]
    ggml_tensor * ff1_lin1_b = nullptr;  // [d_ff], nullable
    ggml_tensor * ff1_lin2_w = nullptr;  // [d_ff, d_model]
    ggml_tensor * ff1_lin2_b = nullptr;  // [d_model], nullable

    // Self-attention with relative positional encoding.
    ggml_tensor * norm_attn_w = nullptr;  // [d_model]
    ggml_tensor * norm_attn_b = nullptr;  // [d_model]
    ggml_tensor * attn_q_w    = nullptr;  // [d_model, d_model]
    ggml_tensor * attn_q_b    = nullptr;  // [d_model], nullable
    ggml_tensor * attn_k_w    = nullptr;  // [d_model, d_model]
    ggml_tensor * attn_k_b    = nullptr;  // [d_model], nullable
    ggml_tensor * attn_v_w    = nullptr;  // [d_model, d_model]
    ggml_tensor * attn_v_b    = nullptr;  // [d_model], nullable
    ggml_tensor * attn_out_w  = nullptr;  // [d_model, d_model]
    ggml_tensor * attn_out_b  = nullptr;  // [d_model], nullable
    ggml_tensor * attn_pos_w  = nullptr;  // [d_model, d_model] linear_pos (no bias)
    ggml_tensor * attn_pos_u  = nullptr;  // [n_head, head_dim]
    ggml_tensor * attn_pos_v  = nullptr;  // [n_head, head_dim]

    // Convolution module. Pointwise conv1 is 2*d_model for GLU.
    ggml_tensor * norm_conv_w         = nullptr;  // [d_model]
    ggml_tensor * norm_conv_b         = nullptr;  // [d_model]
    ggml_tensor * conv_pw1_w          = nullptr;  // [1, d_model, 2*d_model]
    ggml_tensor * conv_pw1_b          = nullptr;  // [2*d_model], nullable
    ggml_tensor * conv_dw_w           = nullptr;  // [k, 1, d_model]
    ggml_tensor * conv_dw_b           = nullptr;  // [d_model], nullable
    ggml_tensor * conv_pw2_w          = nullptr;  // [1, d_model, d_model]
    ggml_tensor * conv_pw2_b          = nullptr;  // [d_model], nullable
    // Post-depthwise normalisation. Exactly one pair is populated per block
    // per BlockParams::conv_norm_type: BatchNorm (conv_bn_fused_*, fused at
    // load) or LayerNorm (conv_ln_*, per-channel mean/std at inference).
    ggml_tensor * conv_bn_fused_scale = nullptr;  // [d_model]
    ggml_tensor * conv_bn_fused_bias  = nullptr;  // [d_model]
    ggml_tensor * conv_ln_w           = nullptr;  // [d_model]
    ggml_tensor * conv_ln_b           = nullptr;  // [d_model]

    // Macaron FF2.
    ggml_tensor * norm_ff2_w = nullptr;
    ggml_tensor * norm_ff2_b = nullptr;
    ggml_tensor * ff2_lin1_w = nullptr;
    ggml_tensor * ff2_lin1_b = nullptr;  // nullable
    ggml_tensor * ff2_lin2_w = nullptr;
    ggml_tensor * ff2_lin2_b = nullptr;  // nullable

    // Final per-block layer norm.
    ggml_tensor * norm_out_w = nullptr;
    ggml_tensor * norm_out_b = nullptr;
};

// Per-family conv dispatch policy. direct_pw is shared (detect_direct_pw is
// the same for every family today). The pre-encode conv0 and depthwise sites
// have separate direct-op controls because their shapes and backend tradeoffs
// differ from the block convolutions. Defaults keep the established im2col
// paths except for direct_pw; families opt into direct pre-encode ops.
struct ConvPolicy {
    bool direct_pw                  = true;
    bool direct_conv0_in_pre_encode = false;
    bool direct_dw_in_block         = false;
    bool direct_dw_in_pre_encode    = false;

    // Causal pre_encode convolutions. NeMo's cache-aware streaming swaps
    // every Conv2d in ConvSubsampling for CausalConv2D, padding
    // (left=k-1, right=stride-1) on both spatial axes — for k=3/s=2 that
    // is (left=2, right=1), different from the offline (k-1)/2 symmetric
    // path, and shifts both the freq and time output dims. False on every
    // offline variant; true on nemotron-speech-streaming-en.
    bool causal_pre_encode = false;
};

// Resolve a conv-dispatch toggle from its env overrides. The DIRECT var forces
// true, the NO_DIRECT var forces false (DIRECT wins if both are set); otherwise
// `backend_default` is used. Centralizes the override parsing shared by the
// pointwise and the per-family depthwise dispatch helpers, so the env-var read
// has exactly one definition. `direct_env` / `no_direct_env` are env var names
// (e.g. "TRANSCRIBE_CONV_DIRECT_DW" / "TRANSCRIBE_CONV_NO_DIRECT_DW").
bool resolve_conv_direct(const char * direct_env, const char * no_direct_env, bool backend_default);

// Pointwise conv dispatch auto-detect (same answer for every family today).
// Env overrides: TRANSCRIBE_CONV_NO_DIRECT_PW forces im2col,
// TRANSCRIBE_CONV_DIRECT_PW forces direct. Vulkan defaults to im2col: on
// AMD Renoir, im2col + mul_mat measured ~200 ms/encode faster than direct for
// f32 weights. Metal and CPU prefer direct.
bool detect_direct_pw(const char * backend);

// Scalar knobs for the block forward. All fields required except the
// optional local-attention window (att_context_left / att_context_right).
struct BlockParams {
    int        d_model;
    int        n_head;
    int        conv_kernel;
    ggml_type  kv_type;    // GGML_TYPE_COUNT = auto (f16 if quant, f32 if f32)
    bool       use_flash;  // false -> manual mul_mat + soft_max + mul_mat
    ConvPolicy policy;

    // Local attention window. -1 / -1 = full attention (default).
    // Semantics depend on att_context_style below.
    int att_context_left  = -1;
    int att_context_right = -1;

    // Self-attention context style.
    //
    //   Regular         — when both att_context_left/right are >= 0,
    //     pos_emb is expected to have length (left+right+1) and
    //     rel_pos_mhsa pads matrix_bd with -inf rows so positions
    //     outside the band become -inf in the attention scores after
    //     softmax. Matches NeMo's LocalAttRelPositionalEncoding when
    //     T <= 2W+1 and remains correct (band-restricted) when T
    //     exceeds the window. Default for every offline variant.
    //
    //   ChunkedLimited  — pos_emb stays at the full 2T-1 length and
    //     the caller supplies a precomputed F16 mask of shape
    //     [T_k, T_q, 1, 1] in `attn_chunked_mask` that
    //     rel_pos_mhsa adds onto matrix_bd before flash_attn. Chunk
    //     size = right+1, left_chunks = left/chunk_size; the mask is
    //     0 inside the allowed [q_chunk - left_chunks, q_chunk] band
    //     and -INF elsewhere. NeMo cache-aware streaming.
    enum class AttContextStyle { Regular, ChunkedLimited };
    AttContextStyle att_context_style = AttContextStyle::Regular;

    // Optional precomputed mask for ChunkedLimited. The caller builds
    // this as a graph input shape [T_k, T_q, 1, 1] F16 (broadcasts
    // across heads) and uploads the host-computed pattern after the
    // compute buffer is allocated. Ignored for AttContextStyle::Regular.
    ggml_tensor * attn_chunked_mask = nullptr;

    // Optional key-padding mask for variable-length offline batching.
    // Shape [T_k, 1, 1, B] F32 (0 on real keys, -INF on padded keys);
    // broadcasts across queries (ne[1]) and heads (ne[2]) and is added
    // onto the attention score matrix in rel_pos_mhsa before flash /
    // soft_max, exactly like attn_chunked_mask. Null for single-shot and
    // same-length batches. Keeps a real query from attending to another
    // utterance's padded tail frames.
    ggml_tensor * attn_pad_mask = nullptr;

    // Depthwise conv padding. -1 / -1 means "centred (k-1)/2" on both
    // sides (every offline variant). Otherwise the depthwise conv
    // uses (conv_context_left, conv_context_right) directly. NeMo's
    // `conv_context_size = "causal"` emits (kernel-1, 0).
    int conv_context_left  = -1;
    int conv_context_right = -1;

    // Optional post-GLU valid-frame mask for the conv module. Values 1.0
    // for valid frames and 0.0 for padded overhang; multiplied onto x at
    // ne=[T, d_model, B, 1], so it broadcasts over d_model (ne[1]). Shape
    // is [T, 1, 1, 1] for the single-utterance / buffered-streaming case
    // and [T, 1, B, 1] for variable-length offline batching (the B axis
    // lands in ne[2]). NeMo applies pad_mask after pointwise_conv1+GLU and
    // before the depthwise convolution; null keeps the historical path.
    ggml_tensor * conv_pad_mask = nullptr;

    // Conv-module normalisation choice. BatchNorm uses fused scale +
    // bias precomputed at load time (BlockView::conv_bn_fused_*).
    // LayerNorm computes per-channel mean/std at inference and uses
    // BlockView::conv_ln_* as the affine scale/bias.
    enum class ConvNormType { BatchNorm, LayerNorm };
    ConvNormType conv_norm_type = ConvNormType::BatchNorm;

    // Streaming (cache-aware) inputs/outputs. All null in offline mode
    // (existing graph topology). When non-null the block runs the streaming
    // path with a "virtual T" attention window (T_cache + T_q_new):
    //   streaming_channel_in  [d_model, T_cache] cache_last_channel from the
    //     previous chunk, concat-prepended before Q/K/V projection.
    //   streaming_channel_out new cache slot (tail of x_norm, size T_q_new);
    //     the driver rotates it into the persistent buffer after compute.
    //   streaming_time_in     [k-1, d_model] cache_last_time replacing the
    //     depthwise conv's zero left-pad (causal pad is (k-1, 0)).
    //   streaming_time_out    last k-1 post-pw1+GLU frames -> next left pad.
    // pos_emb and attn_chunked_mask must be sized for the virtual T.
    ggml_tensor * streaming_channel_in  = nullptr;
    ggml_tensor * streaming_channel_out = nullptr;
    ggml_tensor * streaming_time_in     = nullptr;
    ggml_tensor * streaming_time_out    = nullptr;

    // When streaming, the number of new encoder frames produced this call
    // (slices the attention output and the cache-out source). Offline: 0.
    int streaming_T_q_new = 0;

    // Optional streaming KV cache (all four set together, or none):
    // persistent [d_model, T_cache] pre-projected keys/values. When set, K/V
    // are projected only for the new frames, the cache is prepended, and the
    // rotated cache is emitted into the *_out tensors. streaming_channel_in/
    // _out are ignored in this mode (same state, recompute-K/V form).
    ggml_tensor * streaming_kv_k_in  = nullptr;
    ggml_tensor * streaming_kv_v_in  = nullptr;
    ggml_tensor * streaming_kv_k_out = nullptr;
    ggml_tensor * streaming_kv_v_out = nullptr;

    // Optional precomputed rel-pos projection [head_dim, pos_len, n_head, 1]
    // from a prior chunk of identical geometry; when non-null rel_pos_mhsa
    // consumes it directly and ignores pos_emb. Streaming-only memoization.
    ggml_tensor * streaming_pos_proj_in = nullptr;

    // Graph the helpers forward-expand their cache-write cpy nodes onto. The
    // streaming cache outputs are SIDE outputs (not reachable from the
    // encoder's `out`), so they must be added explicitly. Required when
    // streaming_*_out is non-null.
    ggml_cgraph * streaming_graph = nullptr;
};

// Low-level helpers (exposed for family-specific hand-built blocks).

// Attach a debug name to a tensor. Safe on nullptr.
ggml_tensor * named(ggml_tensor * t, const char * name);

// LayerNorm with affine: y = gamma * (x - mean) / sqrt(var + eps) + beta.
// `gamma` is required; `beta` may be null for bias-free LN.
ggml_tensor * layer_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * gamma, ggml_tensor * beta);

// Two-linear FeedForward: y = Linear2(SiLU(Linear1(x))). Biases are
// optional — pass nullptr to skip the add.
ggml_tensor * feed_forward(ggml_context * ctx,
                           ggml_tensor *  x,
                           ggml_tensor *  lin1_w,
                           ggml_tensor *  lin1_b,
                           ggml_tensor *  lin2_w,
                           ggml_tensor *  lin2_b);

// Macaron half-residual FF: x + 0.5 * FeedForward(LayerNorm(x)).
ggml_tensor * macaron_ff_residual(ggml_context * ctx,
                                  ggml_tensor *  x,
                                  ggml_tensor *  norm_w,
                                  ggml_tensor *  norm_b,
                                  ggml_tensor *  lin1_w,
                                  ggml_tensor *  lin1_b,
                                  ggml_tensor *  lin2_w,
                                  ggml_tensor *  lin2_b);

// Shaw / Transformer-XL relative-position skew. Turns a
// [pos_len, T_q, H, 1] score matrix into a [pos_len, T_q, H, 1]
// matrix rotated so column k holds the score for relative offset k.
ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x);

// f32-friendly Conv1D. Use instead of ggml_conv_1d for fp32 kernels on
// Metal (see conv_2d_dw_f32 in conformer.cpp).
ggml_tensor * conv_1d_f32(ggml_context * ctx,
                          ggml_tensor *  kernel,
                          ggml_tensor *  data,
                          int            stride,
                          int            padding,
                          int            dilation);

// f32-friendly 2D depthwise conv. Use instead of ggml_conv_2d_dw /
// ggml_conv_2d_dw_direct for fp32 kernels on Metal (see conformer.cpp).
ggml_tensor * conv_2d_dw_f32(ggml_context * ctx,
                             ggml_tensor *  kernel,
                             ggml_tensor *  data,
                             int            s0,
                             int            s1,
                             int            p0,
                             int            p1,
                             int            d0,
                             int            d1);

// Batch-stable depthwise Conv2D via ggml_conv_2d_dw_direct, with the
// F32 kernel promotion the direct op needs. Prefer this over
// conv_2d_dw_f32 / conv_1d_dw_f32 when single-shot and batched runs must
// be bit-identical, or when the depthwise conv is hot: the im2col forms
// inflate the activation by the kernel width. See conformer.cpp.
ggml_tensor * conv_2d_dw_direct_f32(ggml_context * ctx,
                                    ggml_tensor *  kernel,
                                    ggml_tensor *  data,
                                    int            s0,
                                    int            s1,
                                    int            p0,
                                    int            p1,
                                    int            d0,
                                    int            d1);

// f32-friendly 1D depthwise conv. Same Metal reasoning.
ggml_tensor * conv_1d_dw_f32(ggml_context * ctx,
                             ggml_tensor *  kernel,
                             ggml_tensor *  data,
                             int            stride,
                             int            padding,
                             int            dilation);

// Fused BatchNorm: y = x * scale + bias, with 1-D scale/bias
// reshaped for broadcast across the time axis of [T, C, 1, 1].
// Both scale_1d and bias_1d are REQUIRED (not nullable) — they come
// from the load-time BN fusion and must always be present. Callers
// that need a no-op path should skip the call entirely rather than
// pass null.
ggml_tensor * fused_batch_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale_1d, ggml_tensor * bias_1d);

// Add a 1-D bias [C] to a 4-D conv output [W, H, C, N] via broadcast.
// Nullable: passing bias_1d == nullptr is a no-op that returns
// conv_out unchanged.
ggml_tensor * add_conv_bias(ggml_context * ctx, ggml_tensor * conv_out, ggml_tensor * bias_1d);

// Sub-blocks (exposed for families that hand-build a block).

// Convolution sub-block: pointwise -> GLU -> optional valid-frame mask
// -> depthwise -> BN/LN -> SiLU -> pointwise. Operates on post-LayerNorm
// input; the LN is applied by the caller. Reads
// BlockParams::conv_context_{left,right} (depthwise padding),
// BlockParams::conv_pad_mask, and BlockParams::conv_norm_type
// (BN vs LN).
ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & b, const BlockParams & params);

// Relative-position multi-head self-attention. Flash attention when
// use_flash is true; manual mul_mat + soft_max_ext + mul_mat fallback
// otherwise. Reads attention-context fields from BlockParams:
//
//   Regular        + att_context_{left,right} >= 0 -> local sliding
//     window. pos_emb has length (left+right+1); helper pads matrix_bd
//     with -inf rows so keys outside the band drop out of softmax.
//
//   ChunkedLimited                                  -> caller supplies
//     a precomputed [T_k, T_q] F16 mask in BlockParams::attn_chunked_mask
//     and pos_emb has full 2T-1 length. The mask is added onto
//     matrix_bd before flash_attn / soft_max_ext.
//
// Otherwise (Regular + att_context_* == -1) -> unrestricted attention.
// x_q (optional): query-only activation [d_model, T_q, B] with T_q !=
// x->ne[1]. When non-null, K/V (and the keys of the score matrix) come
// from `x` while queries come from `x_q` — the cache-aware streaming
// case where only the new frames need attention output but the cached
// frames still serve as keys/values. Requirements in this mode:
//   - pos_emb length must be T_q + T_k - 1 with the zero-offset row at
//     index T_k - 1 (row i holds relative offset (T_k - 1) - i)
//   - attn_chunked_mask (if any) must be [T_k, T_q]
//   - the flash path is bypassed (manual mul_mat + soft_max only)
// Null (default) keeps the historical self-attention behavior.
//
// k_full / v_full (optional, set together): pre-projected keys/values
// [d_model, T_k, 1, 1] (bias already applied). When set, the helper
// skips its own K/V projections entirely, queries come from `x`, and
// the same rectangular requirements as x_q apply. Streaming KV-cache
// path; mutually exclusive with x_q.
ggml_tensor * rel_pos_mhsa(ggml_context *      ctx,
                           ggml_tensor *       x,
                           ggml_tensor *       pos_emb,
                           const BlockView &   b,
                           const BlockParams & params,
                           ggml_tensor *       x_q    = nullptr,
                           ggml_tensor *       k_full = nullptr,
                           ggml_tensor *       v_full = nullptr);

// Top-level block + pre-encode.

// Optional per-block observer hook. `build_conformer_block` calls `on_point`
// at five canonical residual / post-LN points (compile-time `tag`):
//   "after_ff1", "after_attn", "after_conv", "after_ff2", "out".
// Callers attach a ggml name, stash the tensor, and/or mark it for graph
// output. Pointer-of-struct + user pointer (not std::function) keeps the
// header C-ABI-friendly. Pass nullptr to skip observation.
struct BlockObserver {
    void (*on_point)(void * user, const char * tag, ggml_tensor * t) = nullptr;
    void * user                                                      = nullptr;
};

// Full Conformer block forward: FF1 -> attn -> conv -> FF2 -> LN_out.
// `pos_emb` is shared across every block in an encoder — the caller
// computes it once and threads it through. `obs` may be null; when
// non-null its `on_point` callback fires at each of the five
// canonical points listed on BlockObserver.
ggml_tensor * build_conformer_block(ggml_context *        ctx,
                                    ggml_tensor *         x,
                                    ggml_tensor *         pos_emb,
                                    const BlockView &     b,
                                    const BlockParams &   params,
                                    const BlockObserver * obs = nullptr);

// Per-stage valid-frame masks for variable-length offline batching of the
// pre_encode conv stem. Zero-padding to a common length is NOT transparent
// through this stack: convs carry bias + ReLU, so the padded region becomes
// non-zero after the first conv and then leaks into each utterance's last
// VALID frame via the next stride-2 conv's receptive field (flips trailing
// CTC tokens). NeMo's masked subsampling zeros each conv intermediate's
// padded time region after every ReLU. When PreEncodeValidMasks is non-null,
// build_pre_encode allocates three mask graph inputs ne=[1, H_stage, 1, B]
// (one per ReLU stage), applies them, and writes the handles back for the
// driver to fill. Offline (non-causal) pre_encode only.
struct PreEncodeValidMasks {
    ggml_tensor * mask_s1 = nullptr;  // after relu0
    ggml_tensor * mask_s2 = nullptr;  // after relu3
    ggml_tensor * mask_s3 = nullptr;  // after relu6
};

// DwStridingSubsampling pre_encode stack. Returns the final
// [d_model, T_enc, 1, 1] tensor or nullptr on a shape-sanity failure.
// `name_prefix` (if non-null) names every intermediate "<prefix>.conv0"
// etc. Uses pe.out_w->ne[0] as a d_model sanity check; `error_tag` names the
// family in the diagnostic. valid_masks (optional) enables the
// masked-subsampling path above and is filled with the mask input handles.
ggml_tensor * build_pre_encode(ggml_context *        ctx,
                               const PreEncodeView & pe,
                               ggml_tensor *         mel_in,
                               const ConvPolicy &    policy,
                               const char *          name_prefix,
                               const char *          error_tag,
                               PreEncodeValidMasks * valid_masks = nullptr);

}  // namespace transcribe::conformer