Source code
Revision control
Copy as Markdown
Other Tools
#include "encoder.hpp"
#include "subsampling.hpp"
#include "conformer.hpp"
#include "pos_enc.hpp"
#include "ggml_graph.hpp"
#include "backend.hpp"
#include "graph_builder.hpp"
#include "ggml.h"
#include <cassert>
#include <cmath>
#include <cstdlib>
#include <vector>
namespace pk {
// Decide the self-attention window for an encoder of Tp frames. Returns W>0 to
// use NeMo rel_pos_local_attn [W,W] (banded, O(T*window)); -1 for full attention.
//
// The attention uses the chunk-matmul banded path (build_graph_local_chunked),
// which emits O(1) graph nodes regardless of window, so W can go to NeMo's full
// [128,128] without overflowing the metadata-context budget (backend.cpp
// kGraphSize). (The older pad-and-shift path emitted ~6*(2W+1) nodes/layer,
// which is why this was capped at 32.)
static int local_attn_window(int Tp) {
constexpr int kMaxLocalWindow = 128;
if (const char* e = std::getenv("PARAKEET_ATT_CONTEXT")) {
const int w = std::atoi(e);
if (w <= 0) return -1; // 0 / negative -> force full attention
return w > kMaxLocalWindow ? kMaxLocalWindow : w;
}
// Auto: long audio (~>11 min at 8x subsampling) switches to local attention
// so full O(T^2) attention can't OOM the device.
constexpr int kLocalThreshold = 8192;
return Tp > kLocalThreshold ? kMaxLocalWindow : -1;
}
Encoder::Encoder(const ModelLoader& ml)
: ml_(ml) {
d_model_ = (int)ml.config().d_model;
n_layers_ = (int)ml.config().n_layers;
xscaling_ = ml.config().xscaling;
assert(n_layers_ > 0 && d_model_ > 0);
}
void Encoder::forward(const std::vector<float>& mel, int n_mels, int T,
std::vector<float>& enc_out, int& d_model, int& Tout) const {
std::vector<int> none;
std::vector<std::vector<float>> ignored;
forward_capture(mel, n_mels, T, enc_out, d_model, Tout, none, ignored);
}
void Encoder::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 {
// The WHOLE encoder is ONE ggml graph: subsampling -> xscaling -> N conformer
// layers (each FFN1+attn+conv+FFN2+norm_out in-graph) -> final transpose,
// computed in a SINGLE Backend::compute. (Mel stays plain C++, transposed and
// fed as the input inside the subsampling builder.) Versus the old ~85
// per-utterance graphs, this lets the CPU threadpool parallelise across the
// entire encoder. The per-component graph-BUILDERS (Subsampling/RelPos/
// ConformerLayer::build_graph) are reused verbatim by the unit tests.
layer_outs.assign(capture_layers.size(), {});
// Pos_emb is computed host-side once T' is known (deterministic from T, but
// we read it from the subsampling builder). Storage lives in the pool so it
// outlives Backend::compute. capture_layers map: layer idx -> output slot.
GraphInputPool pool;
Subsampling sub(ml_);
int Tp = 0, valid_len = 0, dm_out = 0;
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0,
[&](ggml_context* ctx) -> ggml_tensor* {
// ---- 1. Subsampling: mel [n_mels,T] -> x [d_model, T'] (+ valid). ----
ggml_tensor* x = sub.build_graph(ctx, mel, n_mels, T, pool, Tp,
valid_len);
dm_out = d_model_;
assert((int)x->ne[0] == d_model_);
// ---- 2. xscaling (gated; off for this model). NeMo scales x. ----
if (xscaling_) {
x = ggml_scale(ctx, x, std::sqrt((float)d_model_));
}
// ---- 3. Positional encoding. Long audio uses NeMo
// rel_pos_local_attn (banded, O(T*window)) so attention can't
// OOM; short audio keeps full attention (NeMo-exact). ----
const int att_w = local_attn_window(Tp);
const bool local = att_w > 0;
const int pos_len = local ? (2 * att_w + 1) : (2 * Tp - 1);
std::vector<float>& pe_host = pool.alloc_f32();
if (local) local_rel_pos_encoding(att_w, att_w, d_model_, pe_host);
else rel_pos_encoding(Tp, d_model_, pe_host); // [pos_len, d_model]
int64_t pe_ne[2] = {d_model_, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pe_host.data(), pe_host.size() * sizeof(float));
// ---- 4. Conformer layer stack (all in-graph). ----
for (int i = 0; i < n_layers_; ++i) {
ConformerLayer layer(ml_, i);
x = layer.build_graph(ctx, x, Tp, pe, pos_len, valid_len, pool,
local ? att_w : -1, local ? att_w : -1);
// Capture requested layer outputs from the SAME graph (row-major
// [T', d_model], matching the layer output orientation).
for (size_t c = 0; c < capture_layers.size(); ++c) {
if (capture_layers[c] == i)
pk::capture_graph_output(x, &layer_outs[c]);
}
}
// ---- 5. Final transpose [d_model, T'] -> [T', d_model] (channels-
// first): ggml transpose+cont gives ne0=T', ne1=d_model ->
// row-major [d_model, T'] = enc_out[c*T' + t]. ----
x = ggml_cont(ctx, ggml_transpose(ctx, x));
return x; // ne [T', d_model] -> row-major [d_model, T']
}, enc_out);
assert(ok && "encoder graph failed");
(void)ok;
d_model = dm_out;
Tout = Tp;
}
void Encoder::forward_batch(const MelBatch& mels,
std::vector<std::vector<float>>& enc_outs,
int& d_model, int& Tout,
std::vector<int>& valid_Tout) const {
// Phase 5: the WHOLE batched encoder is ONE fused ggml graph, mirroring
// forward_capture but at B>1: subsampling -> xscaling -> pos_emb -> N
// conformer layers, all [d_model, T', B]. We return the raw [d_model, Tp, B]
// tensor and do the channels-first transpose host-side while splitting per
// item (see the index mapping below).
GraphInputPool pool;
Subsampling sub(ml_);
int Tp = 0;
std::vector<int> vout;
// Per-item entry valid frames: offline convention is T-1 per clip.
std::vector<int> vin(mels.B);
for (int b = 0; b < mels.B; ++b) vin[b] = mels.valid_T[b] - 1;
std::vector<float> flat; // receives [d_model, Tp, B] (ne0=d_model fastest)
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0,
[&](ggml_context* ctx) -> ggml_tensor* {
// ---- 1. Subsampling (batched): mel -> x [d_model, T', B] (+ valid). ----
ggml_tensor* x = sub.build_graph_batched(ctx, mels.data.data(),
mels.n_mels, mels.T_max, mels.B, pool, Tp, vout, vin);
assert((int)x->ne[0] == d_model_);
// ---- 2. xscaling (gated; off for this model). ----
if (xscaling_) x = ggml_scale(ctx, x, std::sqrt((float)d_model_));
// ---- 3. Positional encoding (local for long audio; see B=1 path). ----
const int att_w = local_attn_window(Tp);
const bool local = att_w > 0;
const int pos_len = local ? (2 * att_w + 1) : (2 * Tp - 1);
std::vector<float>& pe_host = pool.alloc_f32();
if (local) local_rel_pos_encoding(att_w, att_w, d_model_, pe_host);
else rel_pos_encoding(Tp, d_model_, pe_host); // [pos_len, d_model]
int64_t pe_ne[2] = {d_model_, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pe_host.data(), pe_host.size() * sizeof(float));
// ---- 4. Conformer layer stack (all in-graph, shared pe). ----
for (int i = 0; i < n_layers_; ++i) {
ConformerLayer layer(ml_, i);
x = layer.build_graph_batched(ctx, x, Tp, mels.B, pe, pos_len, vout, pool,
local ? att_w : -1, local ? att_w : -1);
}
return x; // [d_model, Tp, B]
}, flat);
assert(ok && "batched encoder graph failed");
(void)ok;
// flat is the [d_model, Tp, B] tensor flattened in ggml order (ne0=d_model
// fastest): element (c, t, b) sits at index ((size_t)b*Tp + t)*d_model_ + c.
// Split per item and transpose to channels-first [d_model, Tp]
// (enc_out[c*Tp + t]), matching what forward()/forward_capture return.
d_model = d_model_;
Tout = Tp;
valid_Tout = vout;
// Each enc_outs[b] is channels-first [d_model, valid_Tout[b]]: compact to the
// per-item non-pad frame count so the row stride equals valid_Tout[b]. The
// fused graph runs every item at the padded width Tp, but the trailing
// (Tp - vout[b]) columns are pad-derived; emitting them would (a) make the
// row stride differ from valid_Tout[b] (decoders index enc_out[c*Tout + t]
// with Tout = valid_Tout[b]) and (b) feed pad frames into the decoder. Both
// corrupt a padded (shorter) item's decode.
enc_outs.assign(mels.B, std::vector<float>());
for (int b = 0; b < mels.B; ++b) {
const int tv = vout[b];
enc_outs[b].resize((size_t)d_model_ * tv);
for (int t = 0; t < tv; ++t)
for (int c = 0; c < d_model_; ++c)
enc_outs[b][(size_t)c * tv + t] =
flat[(((size_t)b * Tp) + t) * d_model_ + c];
}
}
void Encoder::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 {
// Mirror of forward_batch's post-subsampling body (steps 2-4 + per-item
// split), but the subsampled features arrive pre-computed in x0_host and are
// injected as a graph input instead of built via sub.build_graph_batched.
GraphInputPool pool;
std::vector<float> flat; // receives [d_model, Tp, B] (ne0=d_model fastest)
bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0,
[&](ggml_context* ctx) -> ggml_tensor* {
// ---- 1. Inject pre-subsampled features x [d_model, Tp, B]. ----
int64_t x_ne[3] = { d_model_, Tp, B };
ggml_tensor* x = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 3, x_ne,
x0_host.data(), x0_host.size() * sizeof(float));
// ---- 2. xscaling (gated; off for this model). ----
if (xscaling_) x = ggml_scale(ctx, x, std::sqrt((float)d_model_));
// ---- 3. Positional encoding (local for long audio; see B=1 path). ----
const int att_w = local_attn_window(Tp);
const bool local = att_w > 0;
const int pos_len = local ? (2 * att_w + 1) : (2 * Tp - 1);
std::vector<float>& pe_host = pool.alloc_f32();
if (local) local_rel_pos_encoding(att_w, att_w, d_model_, pe_host);
else rel_pos_encoding(Tp, d_model_, pe_host); // [pos_len, d_model]
int64_t pe_ne[2] = {d_model_, pos_len};
ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,
pe_host.data(), pe_host.size() * sizeof(float));
// ---- 4. Conformer layer stack (all in-graph, shared pe). ----
for (int i = 0; i < n_layers_; ++i) {
ConformerLayer layer(ml_, i);
x = layer.build_graph_batched(ctx, x, Tp, B, pe, pos_len, vout, pool,
local ? att_w : -1, local ? att_w : -1);
}
return x; // [d_model, Tp, B]
}, flat);
assert(ok && "tiled post-subsampling encoder graph failed");
(void)ok;
d_model = d_model_;
Tout = Tp;
valid_Tout = vout;
enc_outs.assign(B, std::vector<float>());
for (int b = 0; b < B; ++b) {
const int tv = vout[b];
enc_outs[b].resize((size_t)d_model_ * tv);
for (int t = 0; t < tv; ++t)
for (int c = 0; c < d_model_; ++c)
enc_outs[b][(size_t)c * tv + t] =
flat[(((size_t)b * Tp) + t) * d_model_ + c];
}
}
void Encoder::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 {
Subsampling sub(ml_);
const int B = mels.B;
std::vector<std::vector<float>> sub_b(B); // each [Tp_b, d_model] (t*dm+c)
std::vector<int> Tp_b(B), vout(B);
int Tp_max = 0, dm = 0;
for (int b = 0; b < B; ++b) {
// slice item b's real mel [n_mels, valid_T[b]] out of the padded batch buffer
std::vector<float> mel((size_t)mels.n_mels * mels.valid_T[b]);
for (int m = 0; m < mels.n_mels; ++m)
for (int t = 0; t < mels.valid_T[b]; ++t)
mel[(size_t)m*mels.valid_T[b]+t] =
mels.data[((size_t)b*mels.n_mels+m)*mels.T_max + t];
int Tp=0, vl=0;
sub.forward_tiled(mel, mels.n_mels, mels.valid_T[b], tile_out_frames,
sub_b[b], Tp, dm, vl);
Tp_b[b]=Tp; vout[b]=vl; if (Tp>Tp_max) Tp_max=Tp;
}
// assemble x0 [d_model, Tp_max, B], zero-padded per item.
// sub_b[b] is [Tp_b, d_model] (t*dm+c); x0 wants (c,t,b) at (b*Tp_max+t)*dm+c.
std::vector<float> x0((size_t)dm*Tp_max*B, 0.0f);
for (int b=0;b<B;++b) for (int t=0;t<Tp_b[b];++t) for (int c=0;c<dm;++c)
x0[((size_t)b*Tp_max+t)*dm+c] = sub_b[b][(size_t)t*dm+c];
run_post_subsampling_batch(x0, Tp_max, B, vout, enc_outs, d_model, Tout, valid_Tout);
}
} // namespace pk