Source code
Revision control
Copy as Markdown
Other Tools
#pragma once
#include "model_loader.hpp"
#include <vector>
namespace pk {
// Full FastConformer encoder (NeMo ConformerEncoder.forward / forward_internal):
//
// mel [n_mels, T] -> Subsampling (dw_striding รท8) -> [T', d_model]
// -> (if xscaling) x *= sqrt(d_model)
// -> RelPositionalEncoding produces pos_emb [2T'-1, d_model]
// -> for i in 0..n_layers-1: ConformerLayer(i).forward(...)
// -> transpose to [d_model, T'] (NeMo returns [B, d_model, T'])
//
// The `valid_len` (number of non-pad output frames) is derived from Subsampling
// and threaded into every ConformerLayer (attention + conv pad masking).
// A batch of clips' mel features, pre-stacked to T_max.
struct MelBatch {
std::vector<float> data; // contiguous [B][n_mels][T_max], row-major: data[(b*n_mels + m)*T_max + t]
int n_mels = 0;
int T_max = 0; // padded time length (max over clips)
int B = 0;
std::vector<int> valid_T; // per-item true mel frame count, size B (<= T_max)
};
class Encoder {
public:
explicit Encoder(const ModelLoader& ml);
// mel: row-major [n_mels, T] (feat-major inner = T), i.e. mel[m*T + t].
// enc_out: row-major [d_model, Tout] (channels-first, time fastest), matching
// the baseline `encoder_out` orientation: enc_out[c*Tout + t].
void forward(const std::vector<float>& mel, int n_mels, int T,
std::vector<float>& enc_out, int& d_model, int& Tout) const;
// Batched encoder. Runs all B clips through ONE fused ggml graph
// (subsampling -> conformer stack -> output). Returns per-item encoder
// outputs, each row-major [d_model, valid_Tout[b]] (channels-first), the same
// orientation as forward(): each enc_outs[b] is compacted to that item's
// non-pad frame count so its row stride equals valid_Tout[b] (pad-derived
// trailing frames are dropped). `d_model` and `Tout` (the padded T') are
// filled; `valid_Tout` holds each item's non-pad output-frame count (size B).
// For B=1 this reduces to forward().
void forward_batch(const MelBatch& mels,
std::vector<std::vector<float>>& enc_outs,
int& d_model, int& Tout,
std::vector<int>& valid_Tout) const;
// Long-audio tiled variant of forward_batch: subsampling is done per item via
// Subsampling::forward_tiled (which bounds intermediate tensor size so the
// >2^31-element conv tensors that crash long clips never materialise), then
// the pre-subsampled features are fed into the existing post-subsampling graph
// (xscaling + positional encoding + conformer stack). Same output contract as
// forward_batch. `tile_out_frames` = subsampled output frames per tile.
void forward_batch_tiled(const MelBatch& mels,
std::vector<std::vector<float>>& enc_outs,
int& d_model, int& Tout,
std::vector<int>& valid_Tout,
int tile_out_frames) const;
// Same as forward(), but also captures the per-layer outputs at indices
// `capture_layers` (each row-major [T', d_model]) into `layer_outs` (parallel
// to capture_layers). Used by the parity test to localize divergence.
void forward_capture(const std::vector<float>& mel, int n_mels, int T,
std::vector<float>& enc_out, int& d_model, int& Tout,
const std::vector<int>& capture_layers,
std::vector<std::vector<float>>& layer_outs) const;
private:
// Mirrors forward_batch's post-subsampling body (xscaling + positional
// encoding + conformer stack + per-item channels-first split) for the tiled
// path, but takes the already-subsampled features `x0_host` ([d_model, Tp, B]
// in ggml order, element (c,t,b) at ((size_t)b*Tp+t)*d_model+c) injected as a
// graph input instead of building them from sub.build_graph_batched. `vout`
// holds the per-item valid output-frame counts. `x0_host` must outlive the
// call (run_graph is synchronous), which it does as a by-reference parameter.
void run_post_subsampling_batch(const std::vector<float>& x0_host,
int Tp, int B, const std::vector<int>& vout,
std::vector<std::vector<float>>& enc_outs, int& d_model, int& Tout,
std::vector<int>& valid_Tout) const;
const ModelLoader& ml_;
int d_model_;
int n_layers_;
bool xscaling_;
};
} // namespace pk