const ParakeetConfig& c = ml.config();

d_model_ = (int)c.d_model;

n_layers_ = (int)c.n_layers;

n_heads_ = (int)c.n_heads;

conv_kernel_ = (int)c.conv_kernel;

xscaling_ = c.xscaling;

assert(n_heads_ > 0 && d_model_ % n_heads_ == 0);

d_head_ = d_model_ / n_heads_;

left_pad_ = conv_kernel_ - 1; // causal depthwise conv left-context (8)

assert(c.streaming.present && "StreamingEncoder requires a streaming model");

const StreamingCfg& s = c.streaming;

// chunk_size/shift_size/pre_encode_cache_size are stored as the [a,b] arrays;

// the cache-aware streaming buffer uses the [0] value for the FIRST chunk and

// the [1] value thereafter (CacheAwareStreamingAudioBuffer.__iter__).

auto pick = [](const std::vector<int32_t>& v, int idx, int dflt) {

return (idx < (int)v.size()) ? (int)v[idx] : dflt;

};

chunk_first_ = pick(s.chunk_size, 0, 9);

chunk_main_ = pick(s.chunk_size, 1, 16);

pre_cache_ = pick(s.pre_encode_cache_size, 1, 9);

drop_extra_ = s.drop_extra_pre_encoded;

last_channel_cache_= s.last_channel_cache_size;

valid_out_len_ = s.valid_out_len;

att_left_ = c.att_context_left;

att_right_ = c.att_context_right;

assert(c.causal_downsampling && "streaming model expects causal subsampling");

assert(c.conv_causal && "streaming model expects causal depthwise conv");

assert(c.att_context_style == "chunked_limited");

reset();

step_ = 0;

clc_len_ = 0;

cache_time_.assign(n_layers_,

std::vector<float>((size_t)left_pad_ * d_model_, 0.0f));

cache_channel_.assign(n_layers_,

std::vector<float>((size_t)last_channel_cache_ * d_model_, 0.0f));

int n_mel_frames, bool is_last,

int& n_valid_out) {

const int D = d_model_;

const int n_mels = (int)ml_.config().n_mels;

assert((int)mel_chunk_frames.size() == n_mels * n_mel_frames);

// ---- 1. Compute Tc / drop host-side (deterministic from the subsampling

// arithmetic) so the fused graph knows the shapes up front. The

// subsampling builder itself runs in-graph below; here we only

// re-derive its output time length (3 stride-2 k=3 causal stages),

// matching Subsampling's internal calc. ----

Subsampling sub(ml_);

int Tsub_calc = n_mel_frames;

{

const int all_paddings = 3; // causal

int t = n_mel_frames;

for (int st = 0; st < 3; ++st) t = (t + all_paddings - 3) / 2 + 1;

Tsub_calc = t;

}

const int drop = (step_ != 0) ? drop_extra_ : 0;

int Tc = Tsub_calc - drop;

if (Tc < 0) Tc = 0;

if (Tc == 0) { // no output this step (e.g. a too-short tail chunk)

n_valid_out = 0;

step_ += 1;

return {};

}

// ---- 2. Positional encoding. NeMo pos_enc(cache_len): input_len = Tc +

// cache_len; pos_emb spans 2*input_len-1 positions. cache_len = 70. ----

const int cache_len = last_channel_cache_;

const int Pn = Tc + cache_len;

const int pos_len = 2 * Pn - 1;

const int clc_len_now = clc_len_; // snapshot for the in-graph mask

const std::string cnt = ml_.config().conv_norm_type;

StreamLayerCapture caps[/*max layers*/64];

assert(n_layers_ <= 64);

// ---- 3. The WHOLE chunk forward is ONE ggml graph: subsampling (with the

// pre-encode-cache window) -> drop -> xscaling -> N streaming conformer

// layers (each reading its conv/attn cache as inputs and emitting the

// updated caches as captured outputs) -> valid output slice. Computed

// in a SINGLE Backend::compute (vs the old ~5 graphs/layer). ----

std::vector<float> out_full; // [Tc, D] row-major (the final layer output)

GraphInputPool pool;

bool ok = pk::run_graph(/*mem_bytes*/0, /*n_threads*/0,

[&](ggml_context* ctx) -> ggml_tensor* {

// Subsampling on the full mel window (already includes pre-encode

// overlap): the WHOLE window is real audio (no center-pad), so the

// entry valid length is n_mel_frames. Output [D, Tsub].

int Tp = 0, valid = 0;

ggml_tensor* sx = sub.build_graph(ctx, mel_chunk_frames, n_mels,

n_mel_frames, pool, Tp, valid,

/*in_valid_frames*/n_mel_frames);

assert(Tp == Tsub_calc);

// Drop drop_extra_pre_encoded leading subsampled frames (NeMo

// forward_internal; only when cache is present == every step after the

// first). View columns [drop, Tp) of sx [D, Tp] -> [D, Tc].

ggml_tensor* x = sx;

if (drop > 0) {

x = ggml_view_2d(ctx, sx, D, Tc, sx->nb[1],

(size_t)drop * sx->nb[1]);

x = ggml_cont(ctx, x);

}

// xscaling (gated; off for this model).

if (xscaling_) {

x = ggml_scale(ctx, x, std::sqrt((float)D));

}

// pos_emb [D, pos_len] fed as a graph input (host-computed once).

std::vector<float>& pe_host = pool.alloc_f32();

rel_pos_encoding(Pn, D, pe_host); // row-major [pos_len, D]

int64_t pe_ne[2] = {D, pos_len};

ggml_tensor* pe = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, pe_ne,

pe_host.data(), pe_host.size() * sizeof(float));

// Conformer layer stack (conv + attention caches threaded in-graph).

const int LP = left_pad_;

for (int i = 0; i < n_layers_; ++i) {

// Cache reads as graph inputs (host data set after alloc). The

// host buffers ARE the live cache vectors, alive across compute.

int64_t ch_ne[2] = {D, cache_len};

ggml_tensor* cache_ch_t = pk::graph_input_tensor(ctx,

GGML_TYPE_F32, 2, ch_ne, cache_channel_[i].data(),

cache_channel_[i].size() * sizeof(float));

int64_t t_ne[2] = {D, LP};

ggml_tensor* cache_t_t = pk::graph_input_tensor(ctx,

GGML_TYPE_F32, 2, t_ne, cache_time_[i].data(),

cache_time_[i].size() * sizeof(float));

x = build_stream_layer(ctx, ml_, i, x, Tc, pe, pos_len,

cache_len, clc_len_now, n_heads_, D,

conv_kernel_, att_left_, att_right_, cnt,

cache_ch_t, cache_t_t, pool, &caps[i]);

}

return x; // [D, Tc] -> row-major [Tc, D]

}, out_full);

assert(ok && "streaming fused encoder graph failed"); (void)ok;

// ---- 4. Commit the captured next-step caches (read back from the graph). ----

for (int i = 0; i < n_layers_; ++i) {

assert((int)caps[i].next_channel.size() == cache_len * D);

assert((int)caps[i].next_time.size() == left_pad_ * D);

cache_channel_[i].swap(caps[i].next_channel);

cache_time_[i].swap(caps[i].next_time);

}

// After the layer stack the attention cache_last_channel_len has grown by Tc.

clc_len_ = std::min(cache_len, clc_len_ + Tc);

step_ += 1;

// ---- 5. Slice valid output. keep_all_outputs == is_last: mid-stream keep

// only valid_out_len frames; final chunk keep all Tc frames. ----

int valid = is_last ? Tc : std::min(valid_out_len_, Tc);

n_valid_out = valid;

std::vector<float> out((size_t)valid * D);

std::memcpy(out.data(), out_full.data(), (size_t)valid * D * sizeof(float));

return out;