From e7c2a95666febb151b439de90431238ec0fcfe78 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 15:14:23 +0800 Subject: [PATCH 01/17] placeholder commit for pr From e3aecf0e2efa36ca4e00fe963daff4f92ad091ce Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 17:49:17 +0800 Subject: [PATCH 02/17] granite 4 mem improvements --- src/arch/granite_nar/encoder.cpp | 134 ++++++++----------------- src/arch/granite_nar/encoder.h | 69 +++++-------- src/arch/granite_nar/granite_nar.h | 6 +- src/arch/granite_nar/model.cpp | 156 +++++++++++++++++++++++------ 4 files changed, 197 insertions(+), 168 deletions(-) diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index 57a3a516..5470ff71 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/src/arch/granite_nar/encoder.cpp @@ -4,10 +4,9 @@ // self-attention, GLU conv module with conv_expansion=2, macaron FFN, // mid-layer self-conditioned CTC bypass). NLE additions: // -// - A second CTC head over a BPE vocab. We emit -// frame-level logits as `enc.ctc_bpe_logits` here; the -// posterior-weighted window pool + greedy decode runs host-side at -// run() time. +// - A second CTC head over a BPE vocab. Its posterior-weighted encoder +// states are projected in bounded chunks after this graph; each chunk +// performs greedy argmax before returning to host. // - All-hidden-states capture: we tap the per-block POST-LN output at // indices specified by hp.enc_layer_indices (e.g. [4, 8, 12, -1] // 1-indexed → block outputs after 3, 7, 11, 15). These are @@ -269,6 +268,16 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } std::vector captures(capture_idx.size(), nullptr); + for (size_t k = 0; k < capture_idx.size(); ++k) { + if (capture_idx[k] == n_layers - 1) { + eb.final_capture_offset = static_cast(k) * d_model; + break; + } + } + if (eb.final_capture_offset < 0) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar encoder: projector captures do not include final layer"); + return eb; + } for (int i = 0; i < n_layers; ++i) { const auto & b = weights.enc_blocks[i]; @@ -392,16 +401,6 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // Frame-level BPE CTC head (the pool happens host-side). - ggml_tensor * ctc_bpe = nullptr; - if (weights.enc_top.ctc_bpe_w != nullptr) { - ctc_bpe = ggml_mul_mat(ctx, weights.enc_top.ctc_bpe_w, x); - ctc_bpe = ggml_add(ctx, ctc_bpe, weights.enc_top.ctc_bpe_b); - named(ctc_bpe, "enc.ctc_bpe_logits"); - eb.ctc_bpe_logits = ctc_bpe; - ggml_set_output(ctc_bpe); - } - // Channel concat of captured layer outputs. Walk captures in order // (capture_idx ordered by hp.enc_layer_indices entry). ggml_tensor * cat = nullptr; @@ -428,9 +427,6 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } ggml_build_forward_expand(eb.graph, eb.cat_out); ggml_build_forward_expand(eb.graph, eb.ctc_logits); - if (eb.ctc_bpe_logits) { - ggml_build_forward_expand(eb.graph, eb.ctc_bpe_logits); - } if (eb.mid_blank_probs) { ggml_build_forward_expand(eb.graph, eb.mid_blank_probs); } @@ -465,82 +461,36 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } -// Host-side BPE CTC pool + greedy decode. - -void compute_bpe_ctc_initial_hypothesis(const std::vector & importance_non_blank, - const std::vector & ctc_bpe_logits, - int n_bpe_vocab, - int T_enc, - int pool_window, - int blank_id, - std::vector & out_token_ids) { - out_token_ids.clear(); - if (T_enc <= 0 || pool_window <= 0 || n_bpe_vocab <= 0) { - return; - } - if (static_cast(importance_non_blank.size()) < T_enc) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar BPE CTC: importance has %zu entries, need >= %d", - importance_non_blank.size(), T_enc); - return; - } - const std::vector & non_blank = importance_non_blank; - - // Pool over consecutive windows of pool_window. Only windows whose - // total non-blank posterior is non-zero contribute. Within a window - // we form a weighted sum of the BPE logits, weighted by per-frame - // non_blank_prob, normalised by the sum of weights. - const int n_windows = (T_enc + pool_window - 1) / pool_window; - std::vector pooled(static_cast(n_windows) * n_bpe_vocab, 0.0f); - std::vector valid(n_windows, 0); - for (int w = 0; w < n_windows; ++w) { - const int t0 = w * pool_window; - const int t1 = std::min(t0 + pool_window, T_enc); - float total = 0.0f; - for (int t = t0; t < t1; ++t) { - total += non_blank[t]; - } - if (total <= 1e-9f) { - continue; // all-blank window — emit a blank, which collapse drops - } - float * dst = pooled.data() + static_cast(w) * n_bpe_vocab; - for (int t = t0; t < t1; ++t) { - const float wt = non_blank[t] / total; - const float * row = ctc_bpe_logits.data() + static_cast(t) * n_bpe_vocab; - for (int v = 0; v < n_bpe_vocab; ++v) { - dst[v] += wt * row[v]; - } - } - valid[w] = 1; - } - - // Greedy + collapse repeats + drop blanks. For windows with no - // valid mass, we emit blank (no-op). - int prev = -1; - out_token_ids.reserve(n_windows); - for (int w = 0; w < n_windows; ++w) { - int argmax = blank_id; - if (valid[w]) { - const float * dst = pooled.data() + static_cast(w) * n_bpe_vocab; - float best = dst[0]; - for (int v = 1; v < n_bpe_vocab; ++v) { - if (dst[v] > best) { - best = dst[v]; - argmax = v; - } - } - } - if (argmax != blank_id && argmax != prev) { - // Two BPE-CTC schemes, distinguished by blank_id alone: - // - blank_id == 0 (bpe_output_dim = vocab_size + 1): channel 0 - // is a synthetic blank, channels 1..N hold the LLM token ids — - // recover the LLM id with `argmax - 1`. - // - blank_id != 0 (bpe_output_dim = vocab_size): channels ARE - // the LLM ids directly (blank is the BOS id). No shift. - const int shift = (blank_id == 0) ? 1 : 0; - out_token_ids.push_back(argmax - shift); - } - prev = argmax; +BpeCtcBuild build_bpe_ctc_graph(ggml_context * ctx, + const GraniteNarWeights & weights, + const GraniteNarHParams & hp, + int n_windows) { + BpeCtcBuild bb{}; + if (ctx == nullptr || weights.enc_top.ctc_bpe_w == nullptr || weights.enc_top.ctc_bpe_b == nullptr || + hp.enc_hidden <= 0 || hp.enc_bpe_pool_window <= 0 || n_windows <= 0) { + return bb; + } + + bb.n_windows = n_windows; + + bb.hidden_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.enc_hidden, n_windows); + named(bb.hidden_in, "enc.ctc_bpe.hidden_in"); + ggml_set_input(bb.hidden_in); + + ggml_tensor * logits = ggml_mul_mat(ctx, weights.enc_top.ctc_bpe_w, bb.hidden_in); + logits = ggml_add(ctx, logits, weights.enc_top.ctc_bpe_b); + + bb.token_ids = ggml_argmax(ctx, logits); + named(bb.token_ids, "enc.ctc_bpe.token_ids"); + ggml_set_output(bb.token_ids); + + bb.graph = ggml_new_graph_custom(ctx, /*size=*/1024, /*grads=*/false); + if (bb.graph == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar BPE CTC: ggml_new_graph_custom failed"); + return {}; } + ggml_build_forward_expand(bb.graph, bb.token_ids); + return bb; } } // namespace transcribe::granite_nar diff --git a/src/arch/granite_nar/encoder.h b/src/arch/granite_nar/encoder.h index 5fc24338..1a24285f 100644 --- a/src/arch/granite_nar/encoder.h +++ b/src/arch/granite_nar/encoder.h @@ -5,12 +5,13 @@ // AR granite encoder (block-local Shaw self-attention, conv_expansion=2 // GLU, mid-layer self-conditioned CTC bypass) plus two NAR-only additions: // -// 1. A BPE CTC head (1024 → 100352) over a posterior-weighted +// 1. A BPE CTC head (1024 -> 100352) over a posterior-weighted // window=4 pool of valid frames. We expose the bypass-step char-CTC -// mid_logits (1024 → 348) as `enc.ctc_logits` — the exact tensor +// mid_logits (1024 -> 348) as `enc.ctc_logits` -- the exact tensor // the reference model computes at self_conditioning_layer for the -// self-conditioning residual. The pooled BPE head is computed -// host-side at run time. +// self-conditioning residual. The very wide BPE head runs in bounded +// chunks after the encoder so frame-level vocabulary logits are never +// retained for the full utterance. // 2. All-hidden-states capture: the projector consumes 4 encoder // hidden states (post-LN, pre-bypass at the chosen layer // boundaries; indices [4, 8, 12, -1] 1-indexed → layer outputs @@ -57,9 +58,7 @@ struct EncoderBuild { ggml_tensor * cat_out = nullptr; // [num_enc_layers * hidden, T_enc] // the projector input ggml_tensor * ctc_logits = nullptr; // [enc_out_dim=348, T_enc] - ggml_tensor * ctc_bpe_logits = nullptr; // [bpe_output_dim, T_enc] - // raw frame-level (no pool) - ggml_tensor * mid_blank_probs = nullptr; // [T_enc] — softmax(mid_ctc)[blank]. + ggml_tensor * mid_blank_probs = nullptr; // [T_enc] -- softmax(mid_ctc)[blank]. // Used host-side as the BPE pool's // importance weight (importance = // 1 - blank_prob_mid). @@ -82,8 +81,9 @@ struct EncoderBuild { ggml_tensor * block_0_post_ff2 = nullptr; } dumps; - int n_blocks_local = 0; - int last_block_rem = 0; + int n_blocks_local = 0; + int last_block_rem = 0; + int64_t final_capture_offset = -1; // channel offset in cat_out }; EncoderBuild build_encoder_graph(ggml_context * ctx, @@ -96,39 +96,22 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, std::vector precompute_pos_rows(int context_size, int max_pos_emb); std::vector precompute_last_block_mask(int context_size, int t_enc_remainder); -// Host-side BPE CTC pool + greedy decode. -// -// Input: -// importance_non_blank [T_enc] host row of (1 - blank_prob) values used -// as per-frame pool weights. Must come from the -// *middle* (self-conditioning) CTC head's softmax, -// not the final head — see modeling_ctc.py. -// ctc_bpe [T_enc, n_bpe_vocab] host-row-major (frame-level BPE -// logits). n_bpe_vocab is bpe_output_dim: 100353 -// (vocab_size + 1) for the old scheme, 100352 (vocab_size) -// for the new scheme. -// pool_window 4 in this family -// blank_id selects the decode scheme (see weights.h enc_bpe_blank_id): -// - 0 (old): channel 0 is a synthetic blank, channels -// 1..N hold LLM ids; emitted id = argmax - 1. -// - 100257 (new, BOS): channels ARE the LLM ids directly, -// blank is the BOS id; emitted id = argmax (no shift). -// Output: -// token_ids initial-hypothesis BPE token ids after greedy + collapse + -// blank removal. Used as the text portion of the NLE LLM -// forward (each token gets an eos slot inserted around it). -// -// Reference: NLE NARDecoder.compute_text_ctc_preds. Per-frame non-blank -// frames are bucketed into windows of pool_window, the BPE logits over each -// window are weighted by the (softmaxed) CTC non-blank posterior, then a -// greedy argmax + collapse-repeats + drop-blanks yields the hypothesis. See -// the implementation in encoder.cpp. -void compute_bpe_ctc_initial_hypothesis(const std::vector & importance_non_blank, - const std::vector & ctc_bpe_logits, - int n_bpe_vocab, - int T_enc, - int pool_window, - int blank_id, - std::vector & out_token_ids); +// Bounded BPE-CTC projection. The caller supplies posterior-weighted encoder +// states, one per pooling window. Linearity makes projecting a weighted hidden +// state equivalent to weighting the projected frame logits. The graph returns +// one vocabulary argmax per window, avoiding a full-utterance [vocab, T_enc] +// tensor. +struct BpeCtcBuild { + ggml_tensor * hidden_in = nullptr; // [enc_hidden, n_windows] + ggml_tensor * token_ids = nullptr; // [n_windows] i32 + ggml_cgraph * graph = nullptr; + + int n_windows = 0; +}; + +BpeCtcBuild build_bpe_ctc_graph(ggml_context * ctx, + const GraniteNarWeights & weights, + const GraniteNarHParams & hp, + int n_windows); } // namespace transcribe::granite_nar diff --git a/src/arch/granite_nar/granite_nar.h b/src/arch/granite_nar/granite_nar.h index 3e0a84bf..e705426d 100644 --- a/src/arch/granite_nar/granite_nar.h +++ b/src/arch/granite_nar/granite_nar.h @@ -58,10 +58,8 @@ struct GraniteNarModel final : public transcribe_model { struct GraniteNarSession final : public transcribe_session { // Encoder output buffered between encode and projector/LM. std::vector mel_buf; - std::vector enc_cat_host; // [T_enc, num_encoder_layers * enc_hidden] - std::vector ctc_logits_host; // [T_enc, output_dim] - std::vector ctc_bpe_logits_host; // [N_valid, bpe_output_dim] flat - std::vector proj_out_host; // [n_audio_tokens, llm_dim] + std::vector enc_cat_host; // [T_enc, num_encoder_layers * enc_hidden] + std::vector proj_out_host; // [n_audio_tokens, llm_dim] int32_t t_enc = 0; int32_t n_audio_tokens = 0; diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index ed6d2d30..e420483c 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -10,9 +10,9 @@ // forward pass — no KV cache, is_causal=False) // // The forward composes in one run() call as: -// PCM -> mel -> 2-frame stack -> encoder graph (emit cat_out, -// ctc_logits, ctc_bpe_logits) -> host CTC pool + greedy decode of -// the BPE head -> initial hypothesis -> add_insertion_slots -> +// PCM -> mel -> 2-frame stack -> encoder graph (emit cat_out and +// ctc importance) -> bounded BPE-CTC projection + greedy decode -> +// initial hypothesis -> add_insertion_slots -> // text_ids -> decoder forward graph (with projector audio embeds // concatenated to the front of inputs_embeds) -> text_logits -> // argmax + collapse + drop EOS -> final transcript. @@ -391,6 +391,118 @@ void apply_thread_count(ggml_backend_sched_t sched, int n_threads) { transcribe::configure_sched_n_threads(sched, n_threads); } +// Keep the 100k-wide BPE projection bounded. At five minutes, projecting all +// T_enc frames at once would materialize a 5.7 GiB F32 tensor. Pooling four +// encoder states before the linear projection is mathematically equivalent to +// pooling their logits, cuts the projection work by 4x, and lets each bounded +// chunk return only token IDs to the host. +constexpr int kBpeCtcWindowsPerChunk = 512; + +transcribe_status compute_bpe_ctc_initial_hypothesis(ggml_backend_sched_t sched, + const GraniteNarWeights & weights, + const GraniteNarHParams & hp, + const std::vector & enc_cat, + int64_t cat_h, + int T_enc, + int64_t final_offset, + const std::vector & non_blank, + std::vector & out_token_ids, + int64_t & compute_us) { + out_token_ids.clear(); + compute_us = 0; + const int pool_window = hp.enc_bpe_pool_window; + if (sched == nullptr || T_enc <= 0 || pool_window <= 0 || hp.enc_hidden <= 0 || final_offset < 0 || + final_offset + hp.enc_hidden > cat_h || enc_cat.size() < static_cast(cat_h) * T_enc || + non_blank.size() < static_cast(T_enc)) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + + const int n_windows = (T_enc + pool_window - 1) / pool_window; + const int chunk_windows = std::min(kBpeCtcWindowsPerChunk, n_windows); + + ggml_init_params ip{}; + ip.mem_size = 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * bpe_ctx = ggml_init(ip); + if (bpe_ctx == nullptr) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar BPE CTC: context allocation failed"); + return TRANSCRIBE_ERR_OOM; + } + + BpeCtcBuild bb = build_bpe_ctc_graph(bpe_ctx, weights, hp, chunk_windows); + if (bb.graph == nullptr || bb.hidden_in == nullptr || bb.token_ids == nullptr) { + ggml_free(bpe_ctx); + return TRANSCRIBE_ERR_GGUF; + } + + ggml_backend_sched_reset(sched); + if (!ggml_backend_sched_alloc_graph(sched, bb.graph)) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar BPE CTC: graph allocation failed -- out of memory"); + ggml_free(bpe_ctx); + return TRANSCRIBE_ERR_OOM; + } + + const int hidden = hp.enc_hidden; + std::vector hidden_chunk(static_cast(hidden) * bb.n_windows, 0.0f); + std::vector chunk_ids(bb.n_windows, hp.enc_bpe_blank_id); + std::vector valid(bb.n_windows, 0); + + int prev = -1; + out_token_ids.reserve(n_windows); + for (int w0 = 0; w0 < n_windows; w0 += chunk_windows) { + std::fill(hidden_chunk.begin(), hidden_chunk.end(), 0.0f); + std::fill(valid.begin(), valid.end(), 0); + + const int actual_windows = std::min(chunk_windows, n_windows - w0); + for (int w = 0; w < actual_windows; ++w) { + const int global_window = w0 + w; + const int t0 = global_window * pool_window; + const int t1 = std::min(t0 + pool_window, T_enc); + float total = 0.0f; + for (int t = t0; t < t1; ++t) { + total += non_blank[static_cast(t)]; + } + if (total <= 1e-9f) { + continue; + } + valid[static_cast(w)] = 1; + float * dst = hidden_chunk.data() + static_cast(w) * hidden; + for (int t = t0; t < t1; ++t) { + const float wt = non_blank[static_cast(t)] / total; + const float * src = enc_cat.data() + static_cast(t) * cat_h + final_offset; + for (int h = 0; h < hidden; ++h) { + dst[h] += wt * src[h]; + } + } + } + + ggml_backend_tensor_set(bb.hidden_in, hidden_chunk.data(), 0, hidden_chunk.size() * sizeof(float)); + + const int64_t t_compute_start = ggml_time_us(); + const ggml_status gs = ggml_backend_sched_graph_compute(sched, bb.graph); + compute_us += ggml_time_us() - t_compute_start; + if (gs != GGML_STATUS_SUCCESS) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite_nar BPE CTC: compute failed (%d)", static_cast(gs)); + ggml_free(bpe_ctx); + return TRANSCRIBE_ERR_GGUF; + } + ggml_backend_tensor_get(bb.token_ids, chunk_ids.data(), 0, chunk_ids.size() * sizeof(int32_t)); + + for (int w = 0; w < actual_windows; ++w) { + const int argmax = valid[static_cast(w)] ? chunk_ids[static_cast(w)] : hp.enc_bpe_blank_id; + if (argmax != hp.enc_bpe_blank_id && argmax != prev) { + const int shift = (hp.enc_bpe_blank_id == 0) ? 1 : 0; + out_token_ids.push_back(argmax - shift); + } + prev = argmax; + } + } + + ggml_free(bpe_ctx); + return TRANSCRIBE_OK; +} + } // namespace transcribe_status run(transcribe_session * ctx_base, @@ -530,18 +642,6 @@ transcribe_status run(transcribe_session * ctx_base, cc->enc_cat_host.resize(static_cast(cat_h) * cat_T); ggml_backend_tensor_get(eb.cat_out, cc->enc_cat_host.data(), 0, cc->enc_cat_host.size() * sizeof(float)); - const int n_ctc_vocab = static_cast(eb.ctc_logits->ne[0]); - cc->ctc_logits_host.resize(static_cast(n_ctc_vocab) * cat_T); - ggml_backend_tensor_get(eb.ctc_logits, cc->ctc_logits_host.data(), 0, cc->ctc_logits_host.size() * sizeof(float)); - - int n_bpe_vocab = 0; - if (eb.ctc_bpe_logits != nullptr) { - n_bpe_vocab = static_cast(eb.ctc_bpe_logits->ne[0]); - cc->ctc_bpe_logits_host.resize(static_cast(n_bpe_vocab) * cat_T); - ggml_backend_tensor_get(eb.ctc_bpe_logits, cc->ctc_bpe_logits_host.data(), 0, - cc->ctc_bpe_logits_host.size() * sizeof(float)); - } - std::vector mid_non_blank; if (eb.mid_blank_probs != nullptr) { std::vector mid_blank(t_enc); @@ -552,21 +652,19 @@ transcribe_status run(transcribe_session * ctx_base, } } - // ggml stores ne[0]=F as the innermost (fastest-varying) axis. For - // a tensor with ne=[F, T_enc], the host buffer layout in memory is - // buf[t * F + f] = element at ggml indices (f, t) - // which is the SAME as `[T_enc, F]` numpy row-major (T outer, F - // inner). No transpose needed: the host-side BPE pool reads - // ctc_bpe_logits_host[t * V + v] - // and gets the v-th logit at frame t correctly. - (void) n_ctc_vocab; - - // Initial BPE hypothesis. + // Initial BPE hypothesis. cat_out is host-row-major [T_enc, cat_h]; + // final_capture_offset identifies the final 1024-wide encoder state used + // by the BPE head inside each row. std::vector hyp_ids; - if (n_bpe_vocab > 0 && !mid_non_blank.empty()) { - compute_bpe_ctc_initial_hypothesis(mid_non_blank, cc->ctc_bpe_logits_host, n_bpe_vocab, t_enc, - cm->hparams.enc_bpe_pool_window, - /*blank_id=*/cm->hparams.enc_bpe_blank_id, hyp_ids); + int64_t bpe_compute_us = 0; + if (!mid_non_blank.empty()) { + const transcribe_status bst = + compute_bpe_ctc_initial_hypothesis(cc->sched, cm->weights, cm->hparams, cc->enc_cat_host, cat_h, t_enc, + eb.final_capture_offset, mid_non_blank, hyp_ids, bpe_compute_us); + if (bst != TRANSCRIBE_OK) { + return bst; + } + cc->t_encode_us += bpe_compute_us; } if (hyp_ids.empty()) { // No tokens — emit empty transcript and return. From 90939bcb667347a1d862a6d46d346c6c08143cb2 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 19:21:16 +0800 Subject: [PATCH 03/17] moonshine slimming --- docs/input-limits.md | 27 ++-- docs/models/moonshine-streaming-tiny.md | 5 + docs/porting/families/moonshine_streaming.md | 8 +- .../hf_cards/moonshine-streaming-tiny.yaml | 5 +- src/arch/moonshine_streaming/capabilities.cpp | 9 +- src/arch/moonshine_streaming/encoder.cpp | 37 +++-- src/arch/moonshine_streaming/encoder.h | 9 +- src/arch/moonshine_streaming/model.cpp | 134 +++++++++++++----- .../moonshine_streaming/moonshine_streaming.h | 9 +- src/transcribe.cpp | 8 +- tests/CMakeLists.txt | 19 ++- .../moonshine_streaming_batch_truncation.cpp | 80 +++++++---- tests/qwen3_asr_batch_truncation.cpp | 4 +- 13 files changed, 217 insertions(+), 137 deletions(-) diff --git a/docs/input-limits.md b/docs/input-limits.md index fb7a6ae0..cb4b6d00 100644 --- a/docs/input-limits.md +++ b/docs/input-limits.md @@ -71,7 +71,7 @@ need and do not have a length gate. | Families | Limit source | Behavior | | --- | --- | --- | -| qwen3_asr, canary_qwen, funasr_nano, granite, granite_nar, voxtral, cohere, canary | decoder context window (`dec_max_position_embeddings` / `dec_max_seq`), or the encoder positional table (`enc_pos_emb_max_len`, for cohere/canary) — all from GGUF | KV cache grows to fit, clamped to the model's true max. Over-length input is **rejected before the decode** (or before the encoder, where the encoder table is the binding limit) with `TRANSCRIBE_ERR_INPUT_TOO_LONG`. | +| qwen3_asr, canary_qwen, funasr_nano, granite, granite_nar, voxtral, cohere, canary, moonshine_streaming | decoder context window (`dec_max_position_embeddings` / `dec_max_seq`), or a learned encoder/adapter positional table (`enc_pos_emb_max_len`, and `adapter.pos_emb` for moonshine_streaming) — all from GGUF | KV cache grows to fit, clamped to the model's true max. Over-length input is **rejected before the decode** (or before the encoder, where a positional table is the binding limit) with `TRANSCRIBE_ERR_INPUT_TOO_LONG`. | These families wrap an LLM-style decoder whose context window (`audio_tokens + prompt + generation`) is the binding constraint. The number of @@ -81,6 +81,12 @@ computes the prefill size *before* running the encoder and rejects an over-length clip immediately — the caller never pays for a compute pass that cannot fit. The rejection goes through the log callback, not raw stderr. +Moonshine Streaming is encoder-decoder rather than audio-LLM, but has the same +hard-gate behavior: its 4096-row learned adapter position table receives one +row per 20 ms encoder frame, imposing an exact **81.92 s** audio limit. This is +also the limit reported by `max_audio_ms`; one-shot, batch, and streaming calls +reject audio past it before an out-of-range embedding lookup can occur. + The one case that cannot be predicted up front is the transcript itself running long enough to exhaust the remaining budget mid-decode (rare — the output would have to be very large for the audio length). There, the run returns the hard @@ -95,7 +101,7 @@ status (the whole-batch call still returns `TRANSCRIBE_OK`). | Families | Window | Behavior | | --- | --- | --- | -| gigaam (~25 s), sensevoice (~30 s), medasr (~400 s), moonshine (output-bound, ~48 s), moonshine_streaming (output-bound, ~17 min) | training / positional window | Any length is accepted; past the window the library emits a `WARN` (degraded accuracy is possible) and proceeds. `max_audio_ms` reports the window as advisory. | +| gigaam (~25 s), sensevoice (~30 s), medasr (~400 s), moonshine (output-bound, ~48 s) | training / positional window | Any length is accepted; past the window the library emits a `WARN` (degraded accuracy is possible) and proceeds. `max_audio_ms` reports the window as advisory. | These families have no hard architectural wall but were trained on a bounded window; beyond it, accuracy degrades rather than failing. The library does not @@ -109,11 +115,6 @@ Moonshine is the honest edge case in this bucket: its cap is on *output* audio length — a dense short clip can hit it too. It is reported via `transcribe_was_truncated()` and a `WARN` (and, offline, the hard `TRANSCRIBE_ERR_OUTPUT_TRUNCATED` status) when the cap is reached. -`moonshine_streaming` has the same output-bound shape with a much larger window -(`dec_max_position_embeddings = 4096`, ≈ 17 min); because it also streams, its -truncation follows the streaming rule below — `stream_finalize` still returns -`TRANSCRIBE_OK` and the truncation surfaces only through -`transcribe_was_truncated()`. ## Context sizing and the `n_ctx` knob @@ -138,9 +139,10 @@ reports the model's default-context ceiling (`n_ctx == 0`); it is not re-derived for a session that narrows `n_ctx`. A session that lowers `n_ctx` may therefore reject audio shorter than the advertised `max_audio_ms`. -Encoder-bound families are different. For cohere and canary, the input-audio -limit is the encoder positional table, while `n_ctx` only bounds the decoder -self-KV / output budget. In those families `transcribe_session_get_limits()` +Encoder-bound families are different. For cohere, canary, and +moonshine_streaming, the input-audio limit is an encoder or adapter positional +table, while `n_ctx` only bounds the decoder self-KV / output budget. In those +families `transcribe_session_get_limits()` reports a smaller `effective_n_ctx` and `max_kv_bytes` when `n_ctx` is lowered, but `effective_max_audio_ms` stays pinned to the encoder input bound. @@ -182,9 +184,8 @@ and has its own terminal-state machine (`transcribe_stream_*`, IDLE/ACTIVE/FINISHED/FAILED), and `stream_feed` / `stream_finalize` return the status of *that step*, not a verdict on the whole transcript. So when a streaming decode reaches its context cap (e.g. `voxtral_realtime` at its -absolute position limit — hours of continuous audio, or `moonshine_streaming` -at its output window), the stream does **not** fail and `stream_finalize` -returns `TRANSCRIBE_OK`; the truncation is surfaced through +absolute position limit — hours of continuous audio), the stream does **not** +fail and `stream_finalize` returns `TRANSCRIBE_OK`; the truncation is surfaced through `transcribe_was_truncated(session)` and a `WARN`. This is deliberate: forcing a stream into a failed terminal state on truncation would discard the committed text the caller has been consuming. A streaming caller that needs to diff --git a/docs/models/moonshine-streaming-tiny.md b/docs/models/moonshine-streaming-tiny.md index 05882c21..ee49c0f3 100644 --- a/docs/models/moonshine-streaming-tiny.md +++ b/docs/models/moonshine-streaming-tiny.md @@ -14,6 +14,11 @@ multilingual capability, and does not emit timestamps. See Useful Sensors' [model card](https://huggingface.co/UsefulSensors/moonshine-streaming-tiny) for training data, intended use, and upstream evaluation methodology. +The learned adapter position table has 4096 rows at one row per 20 ms encoder +frame, giving an exact **81.92-second** input limit. Longer one-shot, batch, or +streaming input returns `TRANSCRIBE_ERR_INPUT_TOO_LONG`; split longer recordings +into utterances before transcription. + Licensed MIT. Ported from upstream commit [`f8e9dfd`](https://huggingface.co/UsefulSensors/moonshine-streaming-tiny/commit/f8e9dfd8c562c257c151a907b7b7f2fe8ff8511a), pinned 2026-05-06. diff --git a/docs/porting/families/moonshine_streaming.md b/docs/porting/families/moonshine_streaming.md index bfdd2ea9..f23ba8f9 100644 --- a/docs/porting/families/moonshine_streaming.md +++ b/docs/porting/families/moonshine_streaming.md @@ -266,8 +266,9 @@ Highlights: 7. `tie_word_embeddings=false` — converter must NOT tie, GGUF must carry an explicit `lm_head` tensor. 8. `pad_token_id=0` (vs moonshine's 2). Tokenizer `vocab_sha256` differs. -9. `max_position_embeddings=4096` (vs moonshine's 194). Decoder KV cache - sizing must accommodate the longer max length. +9. `max_position_embeddings=4096` (vs moonshine's 194) sizes both the decoder + positions and the learned adapter position table. At one adapter row per + 20 ms encoder frame, the table imposes a hard 81.92-second audio limit. ## Capability Validation @@ -332,8 +333,7 @@ projection would. in encoder-frame units (with `frontend_pad = 4` enc frames of conv-stack history beyond the L_total mask context), encode, then on the emit slice `[T_emitted, stable_T)`: - - apply the adapter with absolute pos_ids → append to - `stream_adapter_committed`; + - apply the adapter with absolute pos_ids; - run the cross-KV projection graph → append per-layer K and V to `stream_cross_k_committed[il]` / `stream_cross_v_committed[il]`. 4. `T_emitted = stable_T`. Bump `audio_committed_ms` to match. diff --git a/scripts/hf_cards/moonshine-streaming-tiny.yaml b/scripts/hf_cards/moonshine-streaming-tiny.yaml index d92e2581..c5ac2ab9 100644 --- a/scripts/hf_cards/moonshine-streaming-tiny.yaml +++ b/scripts/hf_cards/moonshine-streaming-tiny.yaml @@ -34,8 +34,9 @@ summary: | English speech-to-text in both one-shot and streaming modes. A 34M-parameter encoder-decoder ASR model designed for streaming use (ergodic encoder + sliding-window attention, 50 Hz time-domain frontend). Takes a 16 kHz mono - WAV and produces a transcript. No translation, no multilingual capability, - no timestamps. + WAV and produces a transcript. The 4096-row learned adapter position table + limits each utterance to 81.92 seconds. No translation, no multilingual + capability, no timestamps. default_quant_index: 2 # Q8_0 diff --git a/src/arch/moonshine_streaming/capabilities.cpp b/src/arch/moonshine_streaming/capabilities.cpp index 5812f181..b381ec62 100644 --- a/src/arch/moonshine_streaming/capabilities.cpp +++ b/src/arch/moonshine_streaming/capabilities.cpp @@ -28,11 +28,10 @@ void apply_family_invariants(transcribe_model & model) { // commits the last partial transcript). caps.supports_streaming = true; - // Streaming latency characteristics (≈240 ms cumulative encoder - // right-context, natural 20 ms emit unit, family-recommended 80 ms - // feed cadence) are documented in the family doc rather than - // advertised as flat caps fields — the model has no inference-time - // latency knob, and supports_streaming above is the generic gate. + // Streaming latency characteristics (approximately 240 ms cumulative + // encoder right-context, natural 20 ms emit unit, family-recommended 80 ms + // feed cadence) are documented in the family doc. The learned adapter + // position table separately imposes the max_audio_ms hard limit. // Cancellation is wired at the per-run + per-feed level. No PNC/ITN // runtime toggle. Whisper-style fallback / long-form / prompt diff --git a/src/arch/moonshine_streaming/encoder.cpp b/src/arch/moonshine_streaming/encoder.cpp index 931ec49e..06bd6d43 100644 --- a/src/arch/moonshine_streaming/encoder.cpp +++ b/src/arch/moonshine_streaming/encoder.cpp @@ -57,7 +57,7 @@ ggml_tensor * asinh_op(ggml_context * ctx, ggml_tensor * z) { // Encoder MHSA without RoPE, with a per-layer sliding-window mask. // // x: [d_model, T_enc] (residual stream dim = enc_d_model) -// mask: [T_enc, T_enc] f32 — uploaded by caller, cast to F16 inside graph +// mask: [T_enc, T_enc] f16 // Returns: [d_model, T_enc] // // q_w/k_w/v_w have ggml ne [d_model, attn_dim]; out_w has ne @@ -65,7 +65,7 @@ ggml_tensor * asinh_op(ggml_context * ctx, ggml_tensor * z) { // smaller than d_model (small/medium); for tiny they are equal. ggml_tensor * mha_encoder_swa(ggml_context * ctx, ggml_tensor * x, - ggml_tensor * mask_f32, + ggml_tensor * mask, ggml_tensor * q_w, ggml_tensor * k_w, ggml_tensor * v_w, @@ -104,18 +104,15 @@ ggml_tensor * mha_encoder_swa(ggml_context * ctx, K = to_attn_layout(K); V = to_attn_layout(V); - // ggml_soft_max_ext / ggml_flash_attn_ext require F16 mask. - ggml_tensor * mask_f16 = ggml_cast(ctx, mask_f32, GGML_TYPE_F16); - ggml_tensor * o; if (use_flash) { - o = ggml_flash_attn_ext(ctx, Q, K, V, mask_f16, scale, 0.0f, 0.0f); + o = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.0f, 0.0f); // FA output: [head_dim_pad, n_heads, T, 1] → [head_dim_pad, T, n_heads, 1] o = ggml_permute(ctx, o, 0, 2, 1, 3); o = ggml_cont(ctx, o); } else { ggml_tensor * kq = ggml_mul_mat(ctx, K, Q); - ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask_f16, scale, 0.0f); + ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, scale, 0.0f); ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, V, 1, 0, 2, 3)); o = ggml_mul_mat(ctx, v_t, kq_soft); // o ne: [head_dim_pad, T, n_heads, 1] @@ -229,17 +226,15 @@ int encoder_t_enc(const MoonshineStreamingHParams & hp, int n_samples) { return T_enc; } -void build_sliding_window_mask(int T_enc, int left_window, int right_window, float * out_mask) { - constexpr float NEG_INF = -std::numeric_limits::infinity(); - // mask[q, k]: q is row (n_q axis), k is col (n_kv axis = innermost). - // ggml mask layout convention: ne0 = n_kv, ne1 = n_q. Row-major - // memory has q as outer index, k as inner. +void build_sliding_window_mask(int T_enc, int left_window, int right_window, ggml_fp16_t * out_mask) { + const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t neg_inf = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); for (int q = 0; q < T_enc; ++q) { for (int k = 0; k < T_enc; ++k) { const int dist = q - k; const bool left_ok = (dist >= 0) && (dist < left_window); const bool right_ok = (dist < 0) && (-dist < right_window); - out_mask[static_cast(q) * T_enc + k] = (left_ok || right_ok) ? 0.0f : NEG_INF; + out_mask[static_cast(q) * T_enc + k] = (left_ok || right_ok) ? zero : neg_inf; } } } @@ -368,9 +363,23 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, transcribe::debug::mark_tensor_for_dump(x); // ---- per-layer sliding-window masks (input tensors) ---- + // Layers with identical geometry share one input. The tiny variant has + // only two unique masks across six layers; retaining six dense T x T + // inputs needlessly dominates long-utterance memory. eb.per_layer_masks.assign(hp.enc_n_layers, nullptr); for (int i = 0; i < hp.enc_n_layers; ++i) { - ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_enc, T_enc); + const int left = hp.enc_sliding_windows[2 * i + 0]; + const int right = hp.enc_sliding_windows[2 * i + 1]; + for (int j = 0; j < i; ++j) { + if (hp.enc_sliding_windows[2 * j + 0] == left && hp.enc_sliding_windows[2 * j + 1] == right) { + eb.per_layer_masks[i] = eb.per_layer_masks[j]; + break; + } + } + if (eb.per_layer_masks[i] != nullptr) { + continue; + } + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, T_enc, T_enc); char mname[64]; std::snprintf(mname, sizeof(mname), "enc.swa_mask.%d", i); named(mask, mname); diff --git a/src/arch/moonshine_streaming/encoder.h b/src/arch/moonshine_streaming/encoder.h index 4fbf6566..6721276d 100644 --- a/src/arch/moonshine_streaming/encoder.h +++ b/src/arch/moonshine_streaming/encoder.h @@ -45,9 +45,8 @@ struct EncoderBuild { // ggml_backend_tensor_set after alloc. ggml_tensor * audio_in = nullptr; - // Per-layer sliding-window attention masks. Each tensor is f32 - // shape [T_enc, T_enc] (n_kv, n_q). Caller uploads from host-built - // mask buffers before computing the graph. + // Per-layer sliding-window attention masks. Each unique window geometry + // has one shared f16 tensor of shape [T_enc, T_enc] (n_kv, n_q). std::vector per_layer_masks; // Output: final encoder hidden state [d_model, T_enc] f32. @@ -89,7 +88,7 @@ EncoderBuild build_encoder_graph(ggml_context * compute_ctx, // (q-k >= 0 && q-k < L) // up to L-1 positions back, including self // || (k-q >= 1 && k-q < R) // up to R-1 positions ahead // -// Caller-provided buffer must be at least T_enc*T_enc floats. -void build_sliding_window_mask(int T_enc, int left_window, int right_window, float * out_mask); +// Caller-provided buffer must contain at least T_enc*T_enc fp16 values. +void build_sliding_window_mask(int T_enc, int left_window, int right_window, ggml_fp16_t * out_mask); } // namespace transcribe::moonshine_streaming diff --git a/src/arch/moonshine_streaming/model.cpp b/src/arch/moonshine_streaming/model.cpp index 524dd31a..378872b6 100644 --- a/src/arch/moonshine_streaming/model.cpp +++ b/src/arch/moonshine_streaming/model.cpp @@ -79,22 +79,15 @@ MoonshineStreamingModel::~MoonshineStreamingModel() { plan.primary_kind = transcribe::BackendKind::Unknown; } -// Input-length contract (see docs/input-limits.md): like moonshine, the wall -// is on *output*, not input. The encoder takes any PCM length; the decode loop -// stops at the decoder position cap (dec_max_position_embeddings, e.g. 4096) -// before EOS, silently truncating. On a cap-exit we set -// transcribe_was_truncated() and WARN (same as qwen3_asr). - -// Advisory transcribe_capabilities::max_audio_ms: the audio the output budget -// (dec_max_position_embeddings tokens) covers at ~4 output tokens/sec. 0 means -// unknown/unbounded. Advisory only — does not reject. -constexpr int k_tokens_per_sec = 4; // rough speech-rate estimate; advisory only - +// The learned adapter position table is shared with the decoder position +// limit. Each encoder frame represents four frontend frames, so it imposes a +// hard audio limit even though encoder attention itself is local. int64_t moonshine_streaming_max_audio_ms(const MoonshineStreamingHParams & hp) { - if (hp.dec_max_position_embeddings <= 0) { + if (hp.dec_max_position_embeddings <= 0 || hp.enc_frame_len <= 0 || hp.fe_sample_rate <= 0) { return 0; } - return static_cast(hp.dec_max_position_embeddings) * 1000 / k_tokens_per_sec; + const int64_t samples_per_encoder_frame = 4LL * hp.enc_frame_len; + return static_cast(hp.dec_max_position_embeddings) * samples_per_encoder_frame * 1000 / hp.fe_sample_rate; } bool kv_cache_init(MoonshineStreamingKvCache & cache, @@ -240,9 +233,11 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } - // Publish the advisory window now that the decoder position cap is known - // (first point it's available after hparams). - m->caps.max_audio_ms = moonshine_streaming_max_audio_ms(m->hparams); + m->caps.max_audio_ms = moonshine_streaming_max_audio_ms(m->hparams); + m->limits.has_context_cap = true; + m->limits.audio_from_caps = true; + m->limits.model_max_ctx = m->hparams.dec_max_position_embeddings; + m->limits.kv_elems_per_ctx_token = static_cast(m->hparams.dec_d_model) * m->hparams.dec_n_layers * 2; gguf_init_params init_params{}; init_params.no_alloc = true; @@ -298,12 +293,11 @@ transcribe_status init_context(transcribe_model * model, auto cc = std::make_unique(); cc->model = model; cc->n_threads = params->n_threads; - cc->kv_type = params->kv_type; + cc->kv_type = params->kv_type == TRANSCRIBE_KV_TYPE_AUTO ? TRANSCRIBE_KV_TYPE_F32 : params->kv_type; + cc->n_ctx = transcribe_session_params_n_ctx(params); - cc->encoder_use_flash = true; // sliding-window mask is uploaded as - // F32 and cast to F16 inside the graph, - // which is the format flash_attn_ext - // expects. Validated under tolerances. + cc->encoder_use_flash = true; // sliding-window masks use the F16 format + // expected by flash_attn_ext. cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); @@ -440,6 +434,39 @@ int decode_generation_budget(const MoonshineStreamingHParams & hp, int T_enc) { return static_cast(audio_samples * k_budget_num / (k_budget_den * k_native_sr_hz) + k_budget_floor); } +int adapter_position_capacity(const MoonshineStreamingModel * cm) { + if (cm == nullptr || cm->weights.adapter.pos_emb_w == nullptr) { + return 0; + } + return std::min(cm->hparams.dec_max_position_embeddings, static_cast(cm->weights.adapter.pos_emb_w->ne[1])); +} + +int padded_encoder_frames(const MoonshineStreamingHParams & hp, int n_samples) { + if (n_samples <= 0 || hp.enc_frame_len <= 0) { + return 0; + } + const int64_t n_frontend_frames = (static_cast(n_samples) + hp.enc_frame_len - 1) / hp.enc_frame_len; + return static_cast((n_frontend_frames + 3) / 4); +} + +bool input_exceeds_adapter_positions(const MoonshineStreamingModel * cm, int n_samples) { + const int capacity = adapter_position_capacity(cm); + return capacity <= 0 || padded_encoder_frames(cm->hparams, n_samples) > capacity; +} + +transcribe_status check_input_position_limit(const MoonshineStreamingModel * cm, int n_samples) { + if (!input_exceeds_adapter_positions(cm, n_samples)) { + return TRANSCRIBE_OK; + } + const int capacity = adapter_position_capacity(cm); + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "moonshine_streaming run: input too long -- %d encoder frames exceed the %d-row adapter position table " + "(max %.2f s)", + padded_encoder_frames(cm->hparams, n_samples), capacity, + moonshine_streaming_max_audio_ms(cm->hparams) / 1000.0); + return TRANSCRIBE_ERR_INPUT_TOO_LONG; +} + // Encoder helper: build the graph for `n_samples` PCM, upload PCM + masks, // compute, and read final-LN output into the caller's host vector. Updates // cc->t_encode_us. emit_dumps fires encoder.* dump points; streaming feed @@ -490,12 +517,16 @@ transcribe_status encode_window_to_host(MoonshineStreamingSession * cc, ggml_backend_tensor_set(eb.audio_in, pcm, 0, static_cast(n_samples) * sizeof(float)); { - std::vector mask_buf(static_cast(T_enc) * T_enc); + std::vector mask_buf(static_cast(T_enc) * T_enc); for (int i = 0; i < hp.enc_n_layers; ++i) { + if (std::find(eb.per_layer_masks.begin(), eb.per_layer_masks.begin() + i, eb.per_layer_masks[i]) != + eb.per_layer_masks.begin() + i) { + continue; + } const int L = hp.enc_sliding_windows[2 * i + 0]; const int R = hp.enc_sliding_windows[2 * i + 1]; build_sliding_window_mask(T_enc, L, R, mask_buf.data()); - ggml_backend_tensor_set(eb.per_layer_masks[i], mask_buf.data(), 0, mask_buf.size() * sizeof(float)); + ggml_backend_tensor_set(eb.per_layer_masks[i], mask_buf.data(), 0, mask_buf.size() * sizeof(ggml_fp16_t)); } } @@ -718,11 +749,17 @@ transcribe_status ensure_kv_cache_for_T(MoonshineStreamingSession * cc, Moonshin resolved_kv = GGML_TYPE_F16; } - if (cc->kv_cache.buffer != nullptr && cc->kv_cache.T_enc != T_enc) { + const int duration_budget = decode_generation_budget(hp, T_enc); + int model_max_ctx = hp.dec_max_position_embeddings > 0 ? hp.dec_max_position_embeddings : 512; + if (cc->n_ctx > 0) { + model_max_ctx = std::min(model_max_ctx, cc->n_ctx); + } + const int n_ctx = duration_budget > 0 ? std::min(duration_budget, model_max_ctx) : model_max_ctx; + + if (cc->kv_cache.buffer != nullptr && (cc->kv_cache.T_enc != T_enc || cc->kv_cache.n_ctx != n_ctx)) { cc->kv_cache.free(); } if (cc->kv_cache.buffer == nullptr) { - const int n_ctx = hp.dec_max_position_embeddings > 0 ? hp.dec_max_position_embeddings : 512; ggml_type cache_type = resolved_kv; if (cache_type == GGML_TYPE_COUNT) { cache_type = GGML_TYPE_F32; @@ -859,8 +896,9 @@ transcribe_status decode_from_kv_cache(MoonshineStreamingSession * cc, // duration budget bounds runaway loops (inputs the model never ends) to a // few dozen tokens; the position cap remains the absolute backstop. A // budget of 0 (unknown frame geometry) falls back to the position cap. - const int dur_budget = decode_generation_budget(hp, T_enc); - const int gen_cap = (dur_budget > 0 && (max_pos <= 0 || dur_budget < max_pos)) ? dur_budget : max_pos; + const int dur_budget = decode_generation_budget(hp, T_enc); + const int natural_cap = (dur_budget > 0 && (max_pos <= 0 || dur_budget < max_pos)) ? dur_budget : max_pos; + const int gen_cap = std::min(natural_cap, cc->kv_cache.n_ctx); std::vector generated_ids; int next_token = -1; @@ -1184,8 +1222,10 @@ transcribe_status run(transcribe_session * session, if (cm == nullptr || cm->plan.scheduler_list.empty()) { return TRANSCRIBE_ERR_INVALID_ARG; } - cc->clear_result(); + if (auto st = check_input_position_limit(cm, n_samples); st != TRANSCRIBE_OK) { + return st; + } const transcribe_status st = run_one_shot_inner(cc, cm, pcm, n_samples, params); if (st != TRANSCRIBE_OK) { return st; @@ -1401,7 +1441,6 @@ transcribe_status stream_begin(transcribe_session * session, cc->stream_pcm_buffer.clear(); cc->stream_pcm_start_sample = 0; - cc->stream_adapter_committed.clear(); cc->stream_cross_k_committed.assign(static_cast(hp.dec_n_layers), std::vector{}); cc->stream_cross_v_committed.assign(static_cast(hp.dec_n_layers), std::vector{}); cc->stream_T_emitted = 0; @@ -1521,8 +1560,6 @@ transcribe_status flush_stable_frames(MoonshineStreamingSession * cc, if (static_cast(adapter_slice.size()) != dec_h * n_emit) { return TRANSCRIBE_ERR_GGUF; } - cc->stream_adapter_committed.insert(cc->stream_adapter_committed.end(), adapter_slice.begin(), adapter_slice.end()); - if (auto st = project_cross_kv_window(cc, cm, adapter_slice.data(), n_emit, cc->stream_cross_k_committed, cc->stream_cross_v_committed); st != TRANSCRIBE_OK) { @@ -1570,6 +1607,18 @@ transcribe_status stream_feed(transcribe_session * session, return TRANSCRIBE_ERR_ABORTED; } + const int64_t prospective_samples = + cc->stream_pcm_start_sample + static_cast(cc->stream_pcm_buffer.size()) + n_samples; + const int position_capacity = adapter_position_capacity(cm); + const int64_t max_samples = static_cast(position_capacity) * cc->stream_samples_per_enc_frame; + if (position_capacity <= 0 || prospective_samples > max_samples) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "moonshine_streaming stream: input too long -- audio exceeds the %d-row adapter position table " + "(max %.2f s)", + position_capacity, moonshine_streaming_max_audio_ms(cm->hparams) / 1000.0); + return TRANSCRIBE_ERR_INPUT_TOO_LONG; + } + cc->stream_pcm_buffer.insert(cc->stream_pcm_buffer.end(), pcm, pcm + n_samples); cc->stream_audio_input_us += samples_to_us(n_samples); @@ -1773,7 +1822,6 @@ void stream_reset(transcribe_session * session) { auto * cc = static_cast(session); cc->stream_pcm_buffer.clear(); cc->stream_pcm_start_sample = 0; - cc->stream_adapter_committed.clear(); for (auto & v : cc->stream_cross_k_committed) { v.clear(); } @@ -1832,10 +1880,17 @@ transcribe_status run_batch(transcribe_session * session, return TRANSCRIBE_ERR_INVALID_ARG; } - const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && - cm->plan.primary_kind != transcribe::BackendKind::Accel && - cm->plan.primary_kind != transcribe::BackendKind::Unknown; - if (n == 1 || !primary_is_gpu || transcribe::debug::enabled()) { + const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && + cm->plan.primary_kind != transcribe::BackendKind::Accel && + cm->plan.primary_kind != transcribe::BackendKind::Unknown; + bool any_input_too_long = false; + for (int i = 0; i < n; ++i) { + if (pcm[i] != nullptr && n_samples[i] > 0 && input_exceeds_adapter_positions(cm, n_samples[i])) { + any_input_too_long = true; + break; + } + } + if (n == 1 || !primary_is_gpu || transcribe::debug::enabled() || any_input_too_long) { return run_batch_serial(cc, pcm, n_samples, n, params); } @@ -1895,8 +1950,11 @@ transcribe_status run_batch(transcribe_session * session, // anyway). The step graph reads only a power-of-two SUB-window of this // that grows with n_past (see below) — reading the full capacity every // step would dominate the decode and make batching a net loss. - const int n_ctx_cap = std::min(max_pos, 2048); - ggml_type kv_type = GGML_TYPE_F32; + int n_ctx_cap = std::min(max_pos, 2048); + if (cc->n_ctx > 0) { + n_ctx_cap = std::min(n_ctx_cap, cc->n_ctx); + } + ggml_type kv_type = GGML_TYPE_F32; if (cc->kv_type == TRANSCRIBE_KV_TYPE_F16) { kv_type = GGML_TYPE_F16; } diff --git a/src/arch/moonshine_streaming/moonshine_streaming.h b/src/arch/moonshine_streaming/moonshine_streaming.h index 8cf76e9c..bead79d6 100644 --- a/src/arch/moonshine_streaming/moonshine_streaming.h +++ b/src/arch/moonshine_streaming/moonshine_streaming.h @@ -119,11 +119,7 @@ struct MoonshineStreamingModel final : public transcribe_model { }; struct MoonshineStreamingSession final : public transcribe_session { - // Host-side mirror of the post-adapter encoder hidden. The adapter - // pos_emb add (and proj when present) is applied once per session; - // this host buffer feeds the cross_kv precompute graph. - std::vector adapter_host; - int enc_T = 0; // T_enc + int enc_T = 0; // T_enc MoonshineStreamingKvCache kv_cache; @@ -133,8 +129,6 @@ struct MoonshineStreamingSession final : public transcribe_session { // ---- incremental streaming state ---- // // Each feed extends host-side committed buffers in lockstep: - // stream_adapter_committed - post-adapter encoder hidden - // [dec_d_model, T_emitted]. // stream_cross_k/v_committed - per decoder layer, [dec_d_model, // T_emitted]; uploaded into the persistent // kv_cache on each partial decode (per-feed, @@ -157,7 +151,6 @@ struct MoonshineStreamingSession final : public transcribe_session { // per-utterance audio + encoder scratch. std::vector stream_pcm_buffer; int64_t stream_pcm_start_sample = 0; - std::vector stream_adapter_committed; std::vector> stream_cross_k_committed; std::vector> stream_cross_v_committed; int32_t stream_T_emitted = 0; diff --git a/src/transcribe.cpp b/src/transcribe.cpp index dfeb5fb4..15f9b68a 100644 --- a/src/transcribe.cpp +++ b/src/transcribe.cpp @@ -2497,10 +2497,10 @@ extern "C" transcribe_status transcribe_session_get_limits(const struct transcri } // max_kv_bytes: worst-case single-utterance KV allocation at the - // effective ceiling, exact for the session's kv_type. The families - // resolve AUTO (and F16) to f16 for the KV cache and use f32 only for - // an explicit F32 request, so the byte size is 4/elem for F32 and - // 2/elem otherwise. This is the ceiling for one utterance, not the + // effective ceiling, exact for the session's resolved kv_type. Most + // families leave AUTO as the F16 default; a family that defaults to + // F32 stores that resolved choice on the session during init. + // This is the ceiling for one utterance, not the // per-run allocation (the cache grows to fit input); transcribe_run_batch // allocates roughly batch_size x this. const int64_t kv_bytes_per_elem = (session->kv_type == TRANSCRIBE_KV_TYPE_F32) ? 4 : 2; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 85b47282..280c70fb 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -854,30 +854,29 @@ if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) SKIP_RETURN_CODE 77) endif() -# Per-utterance OUTPUT_TRUNCATED in transcribe_run_batch (moonshine_streaming). +# Hard adapter-position input limit in one-shot and batch moonshine_streaming. # Skipped (RC 77) when TRANSCRIBE_MOONSHINE_STREAMING_TINY_GGUF is unset. if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) - add_executable(transcribe_moonshine_streaming_batch_truncation + add_executable(transcribe_moonshine_streaming_input_limit moonshine_streaming_batch_truncation.cpp) - target_link_libraries(transcribe_moonshine_streaming_batch_truncation + target_link_libraries(transcribe_moonshine_streaming_input_limit PRIVATE transcribe transcribe-common-example) - target_compile_definitions(transcribe_moonshine_streaming_batch_truncation PRIVATE + target_compile_definitions(transcribe_moonshine_streaming_input_limit PRIVATE "TRANSCRIBE_TEST_SAMPLES_DIR=\"${CMAKE_SOURCE_DIR}/samples\"") - transcribe_apply_warnings(transcribe_moonshine_streaming_batch_truncation) + transcribe_apply_warnings(transcribe_moonshine_streaming_input_limit) - add_test(NAME transcribe_moonshine_streaming_batch_truncation - COMMAND transcribe_moonshine_streaming_batch_truncation) - set_tests_properties(transcribe_moonshine_streaming_batch_truncation PROPERTIES + add_test(NAME transcribe_moonshine_streaming_input_limit + COMMAND transcribe_moonshine_streaming_input_limit) + set_tests_properties(transcribe_moonshine_streaming_input_limit PROPERTIES SKIP_RETURN_CODE 77) endif() # Per-utterance OUTPUT_TRUNCATED in transcribe_run / transcribe_run_batch for a # causal_lm (LLM-decoder) family. Guards the shared src/causal_lm batched step -# loop's truncation detection (the encoder-decoder twin is covered by -# transcribe_moonshine_streaming_batch_truncation). Skipped (RC 77) when +# loop's truncation detection. Skipped (RC 77) when # TRANSCRIBE_QWEN3_ASR_0_6B_GGUF is unset. if(TRANSCRIBE_BUILD_REAL_MODEL_TESTS) add_executable(transcribe_qwen3_asr_batch_truncation diff --git a/tests/moonshine_streaming_batch_truncation.cpp b/tests/moonshine_streaming_batch_truncation.cpp index a2ff6078..4648475d 100644 --- a/tests/moonshine_streaming_batch_truncation.cpp +++ b/tests/moonshine_streaming_batch_truncation.cpp @@ -1,16 +1,10 @@ -// moonshine_streaming_batch_truncation.cpp - real-model gated test that -// transcribe_run_batch reports output truncation PER UTTERANCE. +// moonshine_streaming_batch_truncation.cpp - real-model gated test for the +// learned adapter-position limit in one-shot and batch inference. // -// Moonshine's cap is on output (max_length decode tokens), not input, so a -// long clip runs the decoder into the cap before end-of-stream. In a batch, -// that must surface as a per-utterance TRANSCRIBE_ERR_OUTPUT_TRUNCATED on the -// affected row (with its partial text retained), while a short row that -// finishes normally stays TRANSCRIBE_OK and the whole-batch call still returns -// OK. transcribe_was_truncated() is also set. See docs/input-limits.md. -// -// Batch makeup: -// row 0 = jfk.wav (~11 s) -> completes under the cap -> OK -// row 1 = love-loss.wav (~197 s) -> exceeds the cap -> OUTPUT_TRUNCATED +// The adapter has 4096 position rows and receives one position per 20 ms +// encoder frame, so tiny accepts at most 81.92 seconds. Longer audio must be +// rejected before encoder compute rather than reaching ggml_get_rows with an +// out-of-range index. // // Gating: // - TRANSCRIBE_BUILD_REAL_MODEL_TESTS (CMake, default OFF) builds it. @@ -89,48 +83,72 @@ int main() { transcribe_model_load_params mp; transcribe_model_load_params_init(&mp); - struct transcribe_model * model = nullptr; + transcribe_model * model = nullptr; if (transcribe_model_load_file(model_path, &mp, &model) != TRANSCRIBE_OK) { std::fprintf(stderr, "model load failed: %s\n", model_path); return 1; } + transcribe_capabilities caps; + transcribe_capabilities_init(&caps); + CHECK(transcribe_model_get_capabilities(model, &caps) == TRANSCRIBE_OK); + CHECK_EQ_INT(caps.max_audio_ms, 81920); + transcribe_session_params sp; transcribe_session_params_init(&sp); - struct transcribe_session * s = nullptr; + transcribe_session * s = nullptr; if (transcribe_session_init(model, &sp, &s) != TRANSCRIBE_OK) { - std::fprintf(stderr, "session init failed\n"); + std::fprintf(stderr, "context init failed\n"); transcribe_model_free(model); return 1; } + transcribe_session_limits limits; + transcribe_session_limits_init(&limits); + CHECK(transcribe_session_get_limits(s, &limits) == TRANSCRIBE_OK); + CHECK_EQ_INT(limits.effective_n_ctx, 4096); + CHECK_EQ_INT(limits.effective_max_audio_ms, 81920); + CHECK_EQ_INT(limits.max_kv_bytes, 62914560); + + // One-shot rejects before encoder compute and clears a prior transcript. + CHECK(transcribe_run(s, pcm_short.data(), static_cast(pcm_short.size()), nullptr) == TRANSCRIBE_OK); + CHECK(transcribe_full_text(s) != nullptr && transcribe_full_text(s)[0] != '\0'); + CHECK(transcribe_run(s, pcm_long.data(), static_cast(pcm_long.size()), nullptr) == + TRANSCRIBE_ERR_INPUT_TOO_LONG); + CHECK(transcribe_full_text(s) == nullptr || transcribe_full_text(s)[0] == '\0'); + const float * pcms[2] = { pcm_short.data(), pcm_long.data() }; - const int lens[2] = { (int) pcm_short.size(), (int) pcm_long.size() }; + const int lens[2] = { static_cast(pcm_short.size()), static_cast(pcm_long.size()) }; - // The whole-batch call succeeds even though a row truncates. + // A mixed batch reports the hard input limit per utterance. CHECK(transcribe_run_batch(s, pcms, lens, 2, nullptr) == TRANSCRIBE_OK); CHECK_EQ_INT(transcribe_batch_n_results(s), 2); - - // Row 0 (short) completes; row 1 (long) hits the output cap. CHECK(transcribe_batch_status(s, 0) == TRANSCRIBE_OK); - CHECK(transcribe_batch_status(s, 1) == TRANSCRIBE_ERR_OUTPUT_TRUNCATED); - - // Both rows keep their (partial, for row 1) transcript. - for (int i = 0; i < 2; ++i) { - const char * text = transcribe_batch_full_text(s, i); - CHECK(text != nullptr && text[0] != '\0'); - } - - // The supplemental flag is set whenever any row truncated. - CHECK(transcribe_was_truncated(s) == true); + CHECK(transcribe_batch_status(s, 1) == TRANSCRIBE_ERR_INPUT_TOO_LONG); + const char * short_text = transcribe_batch_full_text(s, 0); + const char * long_text = transcribe_batch_full_text(s, 1); + CHECK(short_text != nullptr && short_text[0] != '\0'); + CHECK(long_text == nullptr || long_text[0] == '\0'); + CHECK(transcribe_was_truncated(s) == false); + + // Streaming rejects a feed that would cross the same table bound. + transcribe_run_params run_params; + transcribe_run_params_init(&run_params); + transcribe_stream_params stream_params; + transcribe_stream_params_init(&stream_params); + CHECK(transcribe_stream_begin(s, &run_params, &stream_params) == TRANSCRIBE_OK); + transcribe_stream_update update; + transcribe_stream_update_init(&update); + CHECK(transcribe_stream_feed(s, pcm_long.data(), static_cast(pcm_long.size()), &update) == + TRANSCRIBE_ERR_INPUT_TOO_LONG); transcribe_session_free(s); transcribe_model_free(model); if (g_failures > 0) { - std::fprintf(stderr, "moonshine_streaming_batch_truncation: %d failures\n", g_failures); + std::fprintf(stderr, "moonshine_streaming_input_limit: %d failures\n", g_failures); return EXIT_FAILURE; } - std::fprintf(stdout, "moonshine_streaming_batch_truncation: ok\n"); + std::fprintf(stdout, "moonshine_streaming_input_limit: ok\n"); return EXIT_SUCCESS; } diff --git a/tests/qwen3_asr_batch_truncation.cpp b/tests/qwen3_asr_batch_truncation.cpp index ae3fb2e0..8a1fa17f 100644 --- a/tests/qwen3_asr_batch_truncation.cpp +++ b/tests/qwen3_asr_batch_truncation.cpp @@ -11,9 +11,7 @@ // transcribe_was_truncated() set) in BOTH paths, while a short clip that // finishes under the budget stays OK and the whole-batch call still returns OK. // -// This is the causal_lm counterpart to moonshine_streaming_batch_truncation -// (which exercises the encoder-decoder batch loop in transcribe-batch-util.cpp). -// It specifically guards the shared src/causal_lm batched step loop's per-row +// This specifically guards the shared src/causal_lm batched step loop's per-row // truncation detection: that loop marks every stopped row `finished` // regardless of WHY it stopped, so truncation must be inferred from the last // sampled token (!= eos), not from `!finished`. A regression there makes a From b6f8adb9797349ac0ba65b6ac013c890a8ae5fe0 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 20:52:55 +0800 Subject: [PATCH 04/17] qwen3 opt --- src/arch/qwen3_asr/encoder.cpp | 181 +++++++++++++++------------------ src/arch/qwen3_asr/encoder.h | 13 +-- src/arch/qwen3_asr/model.cpp | 25 ++--- 3 files changed, 91 insertions(+), 128 deletions(-) diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index c569d56b..9facb297 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -9,6 +9,7 @@ #include "transcribe-debug.h" #include "transcribe-log.h" +#include #include #include #include @@ -70,29 +71,13 @@ std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max return pe; } -std::vector build_cu_seqlens_mask(const EncoderTiming & t, const QwenAsrHParams & hp) { - (void) hp; - const int32_t T = t.T_enc; - // Full attention over the valid aftercnn rows. The measured - // upstream reference (qwen_asr 0.0.6 + transformers eager / sdpa) - // ignores cu_seqlens and runs unmasked bidirectional attention over - // the post-pad-select tensor. vLLM's flash-attn-2 path honors - // cu_seqlens and chunks at window_aftercnn, but its LibriSpeech WER - // is worse than the eager path in our measurements; we follow the - // eager reference so that transcribe.cpp's greedy decode matches - // the per-tensor dumps byte-for-byte after the pad-row trim in - // build_enc_graph. - // - // Zero-filled mask = softmax(scale * QK) with no bias, which is - // identical to soft_max_ext(nullptr) and matches eager semantics. - return std::vector(static_cast(T) * T, 0.0f); -} - // Graph construction namespace { -constexpr float kLayerNormEps = 1e-5f; +constexpr float kLayerNormEps = 1e-5f; +constexpr int kSubsampleChunkBatch = 32; +constexpr int kSubsampleDirectChunks = 64; ggml_tensor * named(ggml_tensor * t, const char * name) { if (t != nullptr && name != nullptr) { @@ -125,9 +110,73 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w, ggml_ return y; } -// One encoder block: pre-LN self-attention (bidirectional, full- -// sequence mask supplied by the caller) + pre-LN GELU FFN. Residuals -// are full 1.0. +ggml_tensor * build_bounded_subsample(ggml_context * ctx, + ggml_tensor * mel_in, + const QwenAsrWeights & weights, + int64_t mel_per_chunk, + int64_t n_mels, + int64_t T_per_chunk, + int64_t ds_h, + int64_t d_model, + int64_t n_chunks) { + // Keep the original single-batch topology for ordinary recordings. Only + // inputs large enough to create excessive im2col workspace are split. + const int64_t chunk_batch = n_chunks <= kSubsampleDirectChunks ? n_chunks : kSubsampleChunkBatch; + + std::vector groups; + groups.reserve(static_cast((n_chunks + chunk_batch - 1) / chunk_batch)); + + for (int64_t chunk0 = 0; chunk0 < n_chunks; chunk0 += chunk_batch) { + const int64_t n_group = std::min(chunk_batch, n_chunks - chunk0); + ggml_tensor * x = ggml_view_4d(ctx, mel_in, mel_per_chunk, n_mels, 1, n_group, mel_in->nb[1], mel_in->nb[2], + mel_in->nb[3], chunk0 * mel_in->nb[3]); + + x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, 2, 2, 1, 1, 1, 1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); + x = ggml_gelu_erf(ctx, x); + x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); + x = ggml_gelu_erf(ctx, x); + x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); + x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); + x = ggml_gelu_erf(ctx, x); + + const int64_t W_ds = x->ne[0]; + const int64_t H_ds = x->ne[1]; + if (W_ds != T_per_chunk) { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "qwen3_asr encoder: post-conv W=%lld does not match " + "per_chunk_aftercnn=%lld", + static_cast(W_ds), static_cast(T_per_chunk)); + return nullptr; + } + + // Reference [B,C,F,T] -> [B,T,C*F] is [W=T,H=F,C,N] -> [F,C,T,N] + // in ggml layout, followed by flattening F*C. + x = ggml_permute(ctx, x, 2, 0, 1, 3); + x = ggml_cont(ctx, x); + x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, n_group); + x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); + x = ggml_reshape_2d(ctx, x, d_model, T_per_chunk * n_group); + groups.push_back(x); + } + + // A balanced tree bounds graph depth and avoids repeatedly copying the + // entire prefix as the number of groups grows. + while (groups.size() > 1) { + std::vector next; + next.reserve((groups.size() + 1) / 2); + for (size_t i = 0; i < groups.size(); i += 2) { + next.push_back(i + 1 < groups.size() ? ggml_concat(ctx, groups[i], groups[i + 1], /*dim=*/1) : groups[i]); + } + groups.swap(next); + } + + return ggml_reshape_3d(ctx, groups.front(), d_model, T_per_chunk, n_chunks); +} + +// One encoder block: pre-LN bidirectional self-attention with an optional +// batch-padding mask, followed by a pre-LN GELU FFN. Residuals are full 1.0. ggml_tensor * build_enc_block(ggml_context * ctx, ggml_tensor * x, ggml_tensor * mask, @@ -243,64 +292,22 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, named(eb.pos_emb_in, "enc.pos_emb.in"); ggml_set_input(eb.pos_emb_in); - eb.mask_in = ggml_new_tensor_2d(ctx, use_flash ? GGML_TYPE_F16 : GGML_TYPE_F32, T_enc, T_enc); - named(eb.mask_in, "enc.attn_mask.in"); - ggml_set_input(eb.mask_in); - // ----- Subsample: 3x Conv2d + GELU + conv_out linear ----- // Reference layout: PyTorch Conv2d on [B, 1, H=n_mels, W=mel_per_chunk]. // For ggml_conv_2d we use [W=mel_per_chunk, H=n_mels, C=1, N=B] — // ggml swaps (W, H) relative to PyTorch, but the 3x3 kernel is // symmetric and stride/pad are (2,2)/(1,1), so the arithmetic is // invariant. Kernels are stored as ne=[KW=3, KH=3, IC, OC]. - ggml_tensor * x = eb.mel_in; - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, - /*s0=*/2, /*s1=*/2, /*p0=*/1, /*p1=*/1, /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); - x = ggml_gelu_erf(ctx, x); - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); - x = ggml_gelu_erf(ctx, x); - - x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); - x = ggml_gelu_erf(ctx, x); - // Now ne = [W=mel_ds, H=n_mels_ds, C=ds_h, N=n_chunks]. - - const int64_t W_ds = x->ne[0]; - const int64_t H_ds = x->ne[1]; - if (W_ds != T_per_chunk) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "qwen3_asr encoder: post-conv W=%lld does not match " - "per_chunk_aftercnn=%lld", - static_cast(W_ds), static_cast(T_per_chunk)); + // Conv im2col workspace scales with its batch axis. Long recordings can + // contain hundreds of independent frontend chunks, so process that axis + // in bounded groups and concatenate only the compact projected outputs. + // Convolution never mixes chunks, making this algebraically identical to + // one large batch while bounding the scheduler allocation. + ggml_tensor * x = + build_bounded_subsample(ctx, eb.mel_in, weights, mel_per_chunk, n_mels, T_per_chunk, ds_h, d_model, n_chunks); + if (x == nullptr) { return eb; } - - // Reshape for the linear projection. Reference does: - // permute(0, 3, 1, 2).contiguous().view(b, t, c*f) - // i.e. [B, C, F_ds, T_ds] -> [B, T_ds, C, F_ds] -> [B, T_ds, C*F_ds]. - // The flat inner axis is (c * F_ds + f). - // - // In ggml ne layout we have [W=T_ds, H=F_ds, C=ds_h, N=B]. To get a - // flat axis of size ds_h * F_ds where c is the slower sub-axis and - // f the faster — matching the reference — we want axes in order - // [f, c, T, B]: ne = [F_ds, ds_h, T_ds, B]. That's a permute(1, 2, 0, 3). - // ggml_permute uses INVERSE semantics vs PyTorch: the i-th argument - // says which NEW axis old axis i goes to (new[a_i] = old[i]). To - // mirror PyTorch's `permute(0,3,1,2)` on [B, C, F, T] (equivalently, - // map [W=13, H=16, C=480, N=11] → [F=16, C=480, T=13, B=11]), we - // need new[0]=old[1], new[1]=old[2], new[2]=old[0], new[3]=old[3], - // which is ggml_permute args (2, 0, 1, 3). - x = ggml_permute(ctx, x, /*a0=*/2, /*a1=*/0, /*a2=*/1, /*a3=*/3); - x = ggml_cont(ctx, x); - x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, n_chunks); - - // Linear conv_out: [ds_h*F_ds, d_model] weight; maps to d_model. - x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); - // ne = [d_model, T_per_chunk, n_chunks] named(x, "enc.subsample.out"); eb.dumps.subsample_out = x; transcribe::debug::mark_tensor_for_dump(x); @@ -346,7 +353,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // ----- 18 encoder blocks ----- const int n_layers = static_cast(weights.enc_blocks.size()); for (int i = 0; i < n_layers; ++i) { - x = build_enc_block(ctx, x, eb.mask_in, weights.enc_blocks[i], static_cast(d_model), + x = build_enc_block(ctx, x, /*mask=*/nullptr, weights.enc_blocks[i], static_cast(d_model), static_cast(n_heads), use_flash); if (i == 0) { named(x, "enc.block.0.out"); @@ -455,34 +462,12 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, ggml_set_input(eb.mask_in); // ----- Subsample: 3x Conv2d + GELU (per-chunk over N = B*n_chunks_max) ----- - ggml_tensor * x = eb.mel_in; - x = ggml_conv_2d(ctx, weights.enc_subsample.conv0_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv0_b); - x = ggml_gelu_erf(ctx, x); - x = ggml_conv_2d(ctx, weights.enc_subsample.conv1_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv1_b); - x = ggml_gelu_erf(ctx, x); - x = ggml_conv_2d(ctx, weights.enc_subsample.conv2_w, x, 2, 2, 1, 1, 1, 1); - x = add_conv_bias(ctx, x, weights.enc_subsample.conv2_b); - x = ggml_gelu_erf(ctx, x); - // ne = [W=T_per_chunk, H=n_mels_ds, C=ds_h, N]. - - const int64_t W_ds = x->ne[0]; - const int64_t H_ds = x->ne[1]; - if (W_ds != T_per_chunk) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "qwen3_asr encoder(batched): post-conv W=%lld != " - "per_chunk_aftercnn=%lld", - static_cast(W_ds), static_cast(T_per_chunk)); + ggml_tensor * x = + build_bounded_subsample(ctx, eb.mel_in, weights, mel_per_chunk, n_mels, T_per_chunk, ds_h, d_model, N); + if (x == nullptr) { return eb; } - // [W=T_ds, H=F_ds, C=ds_h, N] -> [F_ds, ds_h, T_ds, N] (see single-shot). - x = ggml_permute(ctx, x, 2, 0, 1, 3); - x = ggml_cont(ctx, x); - x = ggml_reshape_3d(ctx, x, H_ds * ds_h, T_per_chunk, N); - x = ggml_mul_mat(ctx, weights.enc_subsample.conv_out, x); // [d_model, T_per_chunk, N] - // Add sinusoidal PE (broadcast across the N = B*n_chunks_max axis). x = ggml_add(ctx, x, eb.pos_emb_in); diff --git a/src/arch/qwen3_asr/encoder.h b/src/arch/qwen3_asr/encoder.h index 92db6102..acd231c2 100644 --- a/src/arch/qwen3_asr/encoder.h +++ b/src/arch/qwen3_asr/encoder.h @@ -5,9 +5,8 @@ // qwen_asr.core.transformers_backend.modeling_qwen3_asr // .Qwen3ASRAudioEncoder. // -// Attention: we match the reference's eager full-bidirectional attention over -// the pad-trimmed sequence (see build_cu_seqlens_mask), not vLLM's -// cu_seqlens chunking. +// Attention: full bidirectional attention over the pad-trimmed sequence, +// matching the eager reference rather than vLLM's cu_seqlens chunking. // // Shape conventions (ggml fast-to-slow ne[]): // @@ -23,8 +22,6 @@ // The last chunk's aftercnn trailing pad rows are dropped // in the graph (matches reference's // `padded_embed[padded_mask_after_cnn]` selection). -// attn mask : [T_enc, T_enc] additive. All zeros — full -// bidirectional attention over the valid aftercnn rows. // output : [output_dim, T_enc] #pragma once @@ -64,11 +61,6 @@ EncoderTiming compute_encoder_timing(int32_t n_mel_frames, const QwenAsrHParams // sin(p * inv_ts[k]), the second d_model/2 are cos(p * inv_ts[k]). std::vector build_sinusoid_pe(int32_t d_model, int32_t length, double max_timescale = 10000.0); -// Additive attention bias [T_enc, T_enc]. All zeros — full bidirectional -// attention over the pad-trimmed sequence (the eager baseline, not vLLM's -// cu_seqlens block-diagonal pattern). -std::vector build_cu_seqlens_mask(const EncoderTiming & t, const QwenAsrHParams & hp); - struct EncoderDumps { ggml_tensor * mel_in = nullptr; // graph input ggml_tensor * subsample_out = nullptr; // post conv_out linear, pre-PE @@ -82,7 +74,6 @@ struct EncoderDumps { struct EncoderBuild { ggml_tensor * mel_in = nullptr; // [mel_per_chunk, n_mels, 1, n_chunks] ggml_tensor * pos_emb_in = nullptr; // [d_model, per_chunk_aftercnn] - ggml_tensor * mask_in = nullptr; // [T_enc, T_enc] ggml_tensor * out = nullptr; // [output_dim, T_enc] EncoderDumps dumps{}; ggml_cgraph * graph = nullptr; diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 423f3b38..4965f02e 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -308,14 +308,15 @@ transcribe_status init_context(transcribe_model * model, if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { kv_type = GGML_TYPE_F32; } - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, - /*n_ctx=*/2048, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, - cm->hparams.dec_n_layers, kv_type)) { + const int initial_n_ctx = std::min(1024, qwen3_context_ceiling(cc->n_ctx, cm->hparams)); + if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, initial_n_ctx, cm->hparams.dec_n_kv_heads, + cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, kv_type)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr init_context: KV cache allocation failed " - "(n_ctx=2048, %d kv-heads x %d head-dim x %d layers) — " + "(n_ctx=%d, %d kv-heads x %d head-dim x %d layers) — " "out of memory.", - cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, cm->hparams.dec_n_layers); + initial_n_ctx, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, + cm->hparams.dec_n_layers); return TRANSCRIBE_ERR_OOM; } } @@ -657,20 +658,6 @@ transcribe_status run(transcribe_session * session, ggml_backend_tensor_set(eb.pos_emb_in, pe.data(), 0, pe.size() * sizeof(float)); } - // Attention mask (block-diagonal from cu_seqlens). - { - std::vector mask = build_cu_seqlens_mask(timing, cm->hparams); - if (cc->encoder_use_flash) { - std::vector mask_f16(mask.size()); - for (size_t i = 0; i < mask.size(); ++i) { - mask_f16[i] = ggml_fp32_to_fp16(mask[i]); - } - ggml_backend_tensor_set(eb.mask_in, mask_f16.data(), 0, mask_f16.size() * sizeof(ggml_fp16_t)); - } else { - ggml_backend_tensor_set(eb.mask_in, mask.data(), 0, mask.size() * sizeof(float)); - } - } - transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); From e1e34de0366fc3299bde5a0ea38153980f26c742 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Sun, 13 Sep 2026 21:35:52 +0800 Subject: [PATCH 05/17] gigaam --- src/arch/gigaam/encoder.cpp | 37 +++++++++++++++++++++++- tests/gigaam_workspace_release_smoke.cpp | 3 +- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/arch/gigaam/encoder.cpp b/src/arch/gigaam/encoder.cpp index 2f981ce3..b6d10bfa 100644 --- a/src/arch/gigaam/encoder.cpp +++ b/src/arch/gigaam/encoder.cpp @@ -13,6 +13,7 @@ #include "transcribe-log.h" #include "weights.h" +#include #include #include #include @@ -24,6 +25,9 @@ namespace { namespace conf = transcribe::conformer; +constexpr int64_t kAttentionQueryChunk = 256; +constexpr int64_t kAttentionChunkMinT = 2048; + // Project a GigaamBlock onto the shared BlockView so we can reuse // conf::conv_module. The attention-half fields stay nullptr because we // run a custom rotary path; the conv module fields are the only ones @@ -190,7 +194,7 @@ ggml_tensor * build_rotary_attn(ggml_context * ctx, // Flash attention already writes contiguous // [head_dim, n_head, T, B] data. o = ggml_flash_attn_ext(ctx, q, k, v, /*mask=*/nullptr, scale, 0.0f, 0.0f); - } else { + } else if (T <= kAttentionChunkMinT) { ggml_tensor * kq = ggml_mul_mat(ctx, k, q); // [T_k, T_q, n_head, B] // Additive key-padding mask: ne=[T_k, 1, 1, B] broadcasts over // queries (ne[1]) and heads (ne[2]). -INF/0 are scale-invariant, so @@ -204,6 +208,37 @@ ggml_tensor * build_rotary_attn(ggml_context * ctx, // Manual attention writes [head_dim, T, n_head, B]. Transpose it // to the contiguous layout returned by flash attention. o = ggml_cont(ctx, ggml_permute(ctx, o, 0, 2, 1, 3)); + } else { + // Tile only the query axis. Every query still attends all T keys, so + // this is the same full-context attention with bounded score storage. + ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v, 1, 0, 2, 3)); + std::vector chunks; + chunks.reserve(static_cast((T + kAttentionQueryChunk - 1) / kAttentionQueryChunk)); + for (int64_t q0 = 0; q0 < T; q0 += kAttentionQueryChunk) { + const int64_t n_query = std::min(kAttentionQueryChunk, T - q0); + ggml_tensor * q_chunk = + ggml_view_4d(ctx, q, head_dim, n_query, n_head, Bb, q->nb[1], q->nb[2], q->nb[3], q0 * q->nb[1]); + q_chunk = ggml_cont(ctx, q_chunk); + ggml_tensor * kq = ggml_mul_mat(ctx, k, q_chunk); + if (attn_pad_mask != nullptr) { + kq = ggml_add(ctx, kq, attn_pad_mask); + } + ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr, scale, 0.0f); + ggml_tensor * chunk = ggml_mul_mat(ctx, v_t, kq_soft); + chunk = ggml_cont(ctx, ggml_permute(ctx, chunk, 0, 2, 1, 3)); + chunk = ggml_reshape_3d(ctx, chunk, d_model, n_query, Bb); + chunks.push_back(chunk); + } + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/1) : + chunks[i]); + } + chunks.swap(next); + } + o = chunks.front(); } // Fold [head_dim, n_head] into d_model without another copy. o = ggml_reshape_3d(ctx, o, d_model, T, Bb); diff --git a/tests/gigaam_workspace_release_smoke.cpp b/tests/gigaam_workspace_release_smoke.cpp index 4478d138..7e851b25 100644 --- a/tests/gigaam_workspace_release_smoke.cpp +++ b/tests/gigaam_workspace_release_smoke.cpp @@ -61,7 +61,8 @@ int main() { } const std::vector short_pcm = make_pcm(5.0); - const std::vector long_pcm = make_pcm(45.0); + // Crosses the encoder's long-input query-tiling threshold. + const std::vector long_pcm = make_pcm(90.0); if (transcribe_run(s, short_pcm.data(), static_cast(short_pcm.size()), nullptr) != TRANSCRIBE_OK) { return fail("short run #1"); From 45e03a34d888261c62c103a4e7a5aa8b6f1151a7 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Mon, 14 Sep 2026 00:24:27 +0800 Subject: [PATCH 06/17] multitalker --- src/conformer/conformer.cpp | 127 +++++++++++++++++++++++++++--------- 1 file changed, 96 insertions(+), 31 deletions(-) diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 03d21906..411aef2a 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -13,10 +13,12 @@ #include "transcribe-env.h" #include "transcribe-log.h" +#include #include #include #include #include +#include namespace transcribe::conformer { @@ -306,6 +308,15 @@ ggml_tensor * add_conv_bias(ggml_context * ctx, ggml_tensor * conv_out, ggml_ten return ggml_add(ctx, conv_out, bias_4d); } +ggml_tensor * add_conv_bias_inplace(ggml_context * ctx, ggml_tensor * conv_out, ggml_tensor * bias_1d) { + if (bias_1d == nullptr) { + return conv_out; + } + const int64_t channels = bias_1d->ne[0]; + ggml_tensor * bias_4d = ggml_reshape_4d(ctx, bias_1d, 1, 1, channels, 1); + return ggml_add_inplace(ctx, conv_out, bias_4d); +} + bool resolve_conv_direct(const char * direct_env, const char * no_direct_env, bool backend_default) { if (transcribe::env::flag(direct_env)) { return true; // user override @@ -958,6 +969,46 @@ ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, const char * s return t; } +// Causal depthwise subsampling pads [kernel-1, stride-1] on both spatial +// axes. Slice the output-time axis so only a bounded padded view is live. +ggml_tensor * causal_dw_2d_time_chunked(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data) { + constexpr int64_t kTimeChunk = 512; + constexpr int64_t kStride = 2; + + const int64_t extent = kernel->ne[1]; + const int64_t pad_left = extent - 1; + const int64_t pad_right = kStride - 1; + const int64_t time_out = (data->ne[1] + pad_left + pad_right - extent) / kStride + 1; + kernel = dw_kernel_for_direct(ctx, kernel); + + std::vector chunks; + chunks.reserve(static_cast((time_out + kTimeChunk - 1) / kTimeChunk)); + for (int64_t out0 = 0; out0 < time_out; out0 += kTimeChunk) { + const int64_t n_out = std::min(kTimeChunk, time_out - out0); + const int64_t src0 = out0 * kStride - pad_left; + const int64_t src1 = src0 + (n_out - 1) * kStride + extent; + const int64_t view0 = std::max(0, src0); + const int64_t view1 = std::min(data->ne[1], src1); + const int64_t pad_top = view0 - src0; + const int64_t pad_bot = src1 - view1; + ggml_tensor * input = ggml_view_4d(ctx, data, data->ne[0], view1 - view0, data->ne[2], data->ne[3], data->nb[1], + data->nb[2], data->nb[3], view0 * data->nb[1]); + input = ggml_pad_ext(ctx, input, pad_left, pad_right, pad_top, pad_bot, 0, 0, 0, 0); + chunks.push_back(ggml_conv_2d_dw_direct(ctx, kernel, input, kStride, kStride, + /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1)); + } + + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/1) : chunks[i]); + } + chunks.swap(next); + } + return chunks.front(); +} + } // namespace // Pre-encode subsampling stack. Op order matches NeMo's @@ -1008,9 +1059,15 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, // Causal pre_encode: NeMo's CausalConv2D pads (left=k-1, right=stride-1) // on both spatial axes before the conv (p=0). Offline variants take the // op-side (k-1)/2 symmetric padding instead. - const bool causal_pe = policy.causal_pre_encode; - const int pe_p_op = causal_pe ? 0 : 1; - auto pad_causal = [&](ggml_tensor * t) { + const bool causal_pe = policy.causal_pre_encode; + const int pe_p_op = causal_pe ? 0 : 1; + auto add_pre_encode_bias = [ctx, causal_pe](ggml_tensor * value, ggml_tensor * bias) { + return causal_pe ? add_conv_bias_inplace(ctx, value, bias) : add_conv_bias(ctx, value, bias); + }; + auto pre_encode_relu = [ctx, causal_pe](ggml_tensor * value) { + return causal_pe ? ggml_relu_inplace(ctx, value) : ggml_relu(ctx, value); + }; + auto pad_causal = [&](ggml_tensor * t) { if (!causal_pe) { return t; } @@ -1049,27 +1106,31 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, /*p0=*/pe_p_op, /*p1=*/pe_p_op, /*d0=*/1, /*d1=*/1); } - x = add_conv_bias(ctx, x, pe.conv0_b); + x = add_pre_encode_bias(x, pe.conv0_b); x = name_prefixed(x, name_prefix, "conv0"); - x = ggml_relu(ctx, x); + x = pre_encode_relu(x); x = name_prefixed(x, name_prefix, "relu0"); x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s1 : nullptr, "pre_encode.valid_mask.s1"); // conv2 (depthwise: channels -> channels, groups=channels, k=3 s=2). // im2col path (conv_2d_dw_f32) when direct_dw_in_pre_encode is false. - x = pad_causal(x); - if (policy.direct_dw_in_pre_encode) { - x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv2_w), x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); + if (policy.direct_dw_in_pre_encode && causal_pe) { + x = causal_dw_2d_time_chunked(ctx, pe.conv2_w, x); } else { - x = conv_2d_dw_f32(ctx, pe.conv2_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); + x = pad_causal(x); + if (policy.direct_dw_in_pre_encode) { + x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv2_w), x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = conv_2d_dw_f32(ctx, pe.conv2_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } } - x = add_conv_bias(ctx, x, pe.conv2_b); + x = add_pre_encode_bias(x, pe.conv2_b); x = name_prefixed(x, name_prefix, "conv2"); // conv3 (pointwise: channels -> channels, k=1 s=1 p=0) @@ -1077,35 +1138,39 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, pe.conv3_b); + x = add_pre_encode_bias(x, pe.conv3_b); x = name_prefixed(x, name_prefix, "conv3"); - x = ggml_relu(ctx, x); + x = pre_encode_relu(x); x = name_prefixed(x, name_prefix, "relu3"); x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s2 : nullptr, "pre_encode.valid_mask.s2"); // conv5 (depthwise) -> conv6 (pointwise) -> ReLU - x = pad_causal(x); - if (policy.direct_dw_in_pre_encode) { - x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv5_w), x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); + if (policy.direct_dw_in_pre_encode && causal_pe) { + x = causal_dw_2d_time_chunked(ctx, pe.conv5_w, x); } else { - x = conv_2d_dw_f32(ctx, pe.conv5_w, x, - /*s0=*/2, /*s1=*/2, - /*p0=*/pe_p_op, /*p1=*/pe_p_op, - /*d0=*/1, /*d1=*/1); + x = pad_causal(x); + if (policy.direct_dw_in_pre_encode) { + x = ggml_conv_2d_dw_direct(ctx, dw_kernel_for_direct(ctx, pe.conv5_w), x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } else { + x = conv_2d_dw_f32(ctx, pe.conv5_w, x, + /*s0=*/2, /*s1=*/2, + /*p0=*/pe_p_op, /*p1=*/pe_p_op, + /*d0=*/1, /*d1=*/1); + } } - x = add_conv_bias(ctx, x, pe.conv5_b); + x = add_pre_encode_bias(x, pe.conv5_b); x = name_prefixed(x, name_prefix, "conv5"); x = ggml_conv_2d(ctx, pe.conv6_w, x, /*s0=*/1, /*s1=*/1, /*p0=*/0, /*p1=*/0, /*d0=*/1, /*d1=*/1); - x = add_conv_bias(ctx, x, pe.conv6_b); + x = add_pre_encode_bias(x, pe.conv6_b); x = name_prefixed(x, name_prefix, "conv6"); - x = ggml_relu(ctx, x); + x = pre_encode_relu(x); x = name_prefixed(x, name_prefix, "relu6"); x = apply_valid_mask(x, valid_masks ? &valid_masks->mask_s3 : nullptr, "pre_encode.valid_mask.s3"); From a16aecbebc6115e73a5d8aaf146ef7c950e81465 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Mon, 14 Sep 2026 02:06:01 +0800 Subject: [PATCH 07/17] canary --- src/arch/canary_qwen/encoder.cpp | 8 +++-- src/conformer/conformer.cpp | 55 +++++++++++++++++++++++++++++--- src/conformer/conformer.h | 5 +++ 3 files changed, 61 insertions(+), 7 deletions(-) diff --git a/src/arch/canary_qwen/encoder.cpp b/src/arch/canary_qwen/encoder.cpp index b434a052..397855de 100644 --- a/src/arch/canary_qwen/encoder.cpp +++ b/src/arch/canary_qwen/encoder.cpp @@ -108,9 +108,11 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, bool use_flash, const char * backend_name) { conf::ConvPolicy policy{}; - policy.direct_pw = conf::detect_direct_pw(backend_name); - policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = false; + policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); + policy.direct_dw_in_pre_encode = false; + policy.inplace_pre_encode = true; + policy.pre_encode_dw_time_chunk = 256; EncoderBuild eb{}; diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 411aef2a..704b9b0a 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -969,6 +969,48 @@ ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, const char * s return t; } +// Retain the regular im2col depthwise kernel while bounding its expanded +// output-time axis. Each slice includes its complete source receptive field. +ggml_tensor * regular_dw_2d_time_chunked(ggml_context * ctx, + ggml_tensor * kernel, + ggml_tensor * data, + int64_t time_chunk) { + constexpr int64_t kStride = 2; + + const int64_t extent = kernel->ne[1]; + const int64_t pad = (extent - 1) / 2; + const int64_t time_out = (data->ne[1] + 2 * pad - extent) / kStride + 1; + if (time_out <= time_chunk) { + return conv_2d_dw_f32(ctx, kernel, data, kStride, kStride, pad, pad, /*d0=*/1, /*d1=*/1); + } + + std::vector chunks; + chunks.reserve(static_cast((time_out + time_chunk - 1) / time_chunk)); + for (int64_t out0 = 0; out0 < time_out; out0 += time_chunk) { + const int64_t n_out = std::min(time_chunk, time_out - out0); + const int64_t src0 = out0 * kStride - pad; + const int64_t src1 = src0 + (n_out - 1) * kStride + extent; + const int64_t view0 = std::max(0, src0); + const int64_t view1 = std::min(data->ne[1], src1); + const int64_t pad_top = view0 - src0; + const int64_t pad_bot = src1 - view1; + ggml_tensor * input = ggml_view_4d(ctx, data, data->ne[0], view1 - view0, data->ne[2], data->ne[3], data->nb[1], + data->nb[2], data->nb[3], view0 * data->nb[1]); + input = ggml_pad_ext(ctx, input, 0, 0, pad_top, pad_bot, 0, 0, 0, 0); + chunks.push_back(conv_2d_dw_f32(ctx, kernel, input, kStride, kStride, pad, /*p1=*/0, /*d0=*/1, /*d1=*/1)); + } + + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/1) : chunks[i]); + } + chunks.swap(next); + } + return chunks.front(); +} + // Causal depthwise subsampling pads [kernel-1, stride-1] on both spatial // axes. Slice the output-time axis so only a bounded padded view is live. ggml_tensor * causal_dw_2d_time_chunked(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data) { @@ -1060,12 +1102,13 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, // on both spatial axes before the conv (p=0). Offline variants take the // op-side (k-1)/2 symmetric padding instead. const bool causal_pe = policy.causal_pre_encode; + const bool inplace_pe = causal_pe || policy.inplace_pre_encode; const int pe_p_op = causal_pe ? 0 : 1; - auto add_pre_encode_bias = [ctx, causal_pe](ggml_tensor * value, ggml_tensor * bias) { - return causal_pe ? add_conv_bias_inplace(ctx, value, bias) : add_conv_bias(ctx, value, bias); + auto add_pre_encode_bias = [ctx, inplace_pe](ggml_tensor * value, ggml_tensor * bias) { + return inplace_pe ? add_conv_bias_inplace(ctx, value, bias) : add_conv_bias(ctx, value, bias); }; - auto pre_encode_relu = [ctx, causal_pe](ggml_tensor * value) { - return causal_pe ? ggml_relu_inplace(ctx, value) : ggml_relu(ctx, value); + auto pre_encode_relu = [ctx, inplace_pe](ggml_tensor * value) { + return inplace_pe ? ggml_relu_inplace(ctx, value) : ggml_relu(ctx, value); }; auto pad_causal = [&](ggml_tensor * t) { if (!causal_pe) { @@ -1123,6 +1166,8 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, /*s0=*/2, /*s1=*/2, /*p0=*/pe_p_op, /*p1=*/pe_p_op, /*d0=*/1, /*d1=*/1); + } else if (!causal_pe && policy.pre_encode_dw_time_chunk > 0) { + x = regular_dw_2d_time_chunked(ctx, pe.conv2_w, x, policy.pre_encode_dw_time_chunk); } else { x = conv_2d_dw_f32(ctx, pe.conv2_w, x, /*s0=*/2, /*s1=*/2, @@ -1154,6 +1199,8 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, /*s0=*/2, /*s1=*/2, /*p0=*/pe_p_op, /*p1=*/pe_p_op, /*d0=*/1, /*d1=*/1); + } else if (!causal_pe && policy.pre_encode_dw_time_chunk > 0) { + x = regular_dw_2d_time_chunked(ctx, pe.conv5_w, x, policy.pre_encode_dw_time_chunk); } else { x = conv_2d_dw_f32(ctx, pe.conv5_w, x, /*s0=*/2, /*s1=*/2, diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index 7ed7af9d..3df9afb4 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -112,6 +112,11 @@ struct ConvPolicy { bool direct_dw_in_block = false; bool direct_dw_in_pre_encode = false; + // Optional memory controls for the three-stage subsampler. A positive + // chunk size bounds regular depthwise im2col along output time. + bool inplace_pre_encode = false; + int pre_encode_dw_time_chunk = 0; + // 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 From 916a08e27950613a0f40af5c8576122bebc4fa19 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Mon, 14 Sep 2026 03:41:39 +0800 Subject: [PATCH 08/17] cohere --- src/arch/cohere/encoder.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/arch/cohere/encoder.cpp b/src/arch/cohere/encoder.cpp index 0b0819cc..75ebc721 100644 --- a/src/arch/cohere/encoder.cpp +++ b/src/arch/cohere/encoder.cpp @@ -112,10 +112,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, bool use_flash, const char * backend_name) { conf::ConvPolicy policy{}; - policy.direct_pw = conf::detect_direct_pw(backend_name); - const bool direct_dw = detect_direct_dw(backend_name); - policy.direct_dw_in_block = direct_dw; - policy.direct_dw_in_pre_encode = direct_dw; + policy.direct_pw = conf::detect_direct_pw(backend_name); + const bool direct_dw = detect_direct_dw(backend_name); + policy.direct_dw_in_block = direct_dw; + policy.direct_dw_in_pre_encode = direct_dw; + policy.inplace_pre_encode = !direct_dw; + policy.pre_encode_dw_time_chunk = 256; EncoderBuild eb{}; From c6e27ef82ca05fa52b45cc70e23d32f7999b60b0 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Wed, 16 Sep 2026 22:15:54 +0800 Subject: [PATCH 09/17] granite mem reduction --- src/arch/granite/encoder.cpp | 60 +++++++++++++++++++++++++++++++----- src/arch/granite/encoder.h | 6 ++-- src/arch/granite/model.cpp | 6 ++-- 3 files changed, 61 insertions(+), 11 deletions(-) diff --git a/src/arch/granite/encoder.cpp b/src/arch/granite/encoder.cpp index 1396df66..6b54175e 100644 --- a/src/arch/granite/encoder.cpp +++ b/src/arch/granite/encoder.cpp @@ -223,7 +223,8 @@ ggml_tensor * granite_conv_module(ggml_context * ctx, ggml_tensor * bn_fused_scale, ggml_tensor * bn_fused_bias, int conv_kernel, - int inner_dim) { + int inner_dim, + bool direct_depthwise) { const int64_t d_model = x->ne[0]; const int64_t T = x->ne[1]; @@ -250,11 +251,49 @@ ggml_tensor * granite_conv_module(ggml_context * ctx, // Transpose for depthwise conv: [inner_dim, T] -> [T, inner_dim]. x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); - // Depthwise conv1d: kernel [k, 1, inner_dim], symmetric pad (k-1)/2 - // (= 7 for k=15). Granite's depthwise has no bias. + // Depthwise conv1d: kernel [k, 1, inner_dim], symmetric pad (k-1)/2. + // The im2col implementation expands [T, C] to [k*C, T], which reaches + // gigabytes on long clips. Prefer the direct depthwise operator; keep a + // bounded im2col fallback for backends where the direct op is disabled. const int padding = (conv_kernel - 1) / 2; - x = transcribe::conformer::conv_1d_dw_f32(ctx, b.conv_depthwise_w, x, - /*stride=*/1, /*padding=*/padding, /*dilation=*/1); + if (direct_depthwise) { + ggml_tensor * kernel = ggml_reshape_4d(ctx, b.conv_depthwise_w, conv_kernel, 1, 1, inner_dim); + ggml_tensor * data = ggml_reshape_4d(ctx, x, T, 1, inner_dim, 1); + x = transcribe::conformer::conv_2d_dw_direct_f32(ctx, kernel, data, + /*s0=*/1, /*s1=*/1, + /*p0=*/padding, /*p1=*/0, + /*d0=*/1, /*d1=*/1); + x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[2], x->ne[3]); + } else { + constexpr int64_t kTimeChunk = 256; + std::vector chunks; + chunks.reserve(static_cast((T + kTimeChunk - 1) / kTimeChunk)); + for (int64_t out0 = 0; out0 < T; out0 += kTimeChunk) { + const int64_t n_out = std::min(kTimeChunk, T - out0); + const int64_t src0 = out0 - padding; + const int64_t src1 = src0 + n_out + conv_kernel - 1; + const int64_t view0 = std::max(0, src0); + const int64_t view1 = std::min(T, src1); + const int64_t pad_left = view0 - src0; + const int64_t pad_right = src1 - view1; + ggml_tensor * input = + ggml_view_3d(ctx, x, view1 - view0, inner_dim, 1, x->nb[1], x->nb[2], view0 * x->nb[0]); + input = ggml_pad_ext(ctx, input, pad_left, pad_right, 0, 0, 0, 0, 0, 0); + chunks.push_back(transcribe::conformer::conv_1d_dw_f32(ctx, b.conv_depthwise_w, input, + /*stride=*/1, /*padding=*/0, + /*dilation=*/1)); + } + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/0) : + chunks[i]); + } + chunks.swap(next); + } + x = chunks.front(); + } // Fused BatchNorm: y = x * fused_scale + fused_bias, with 1-D scale // and bias broadcast across the time axis. @@ -285,8 +324,15 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const GraniteWeights & weights, const GraniteHParams & hp, int T_enc, - bool /*use_flash*/) { + bool /*use_flash*/, + const char * backend_name) { EncoderBuild eb{}; + + const bool backend_direct = backend_name != nullptr && (std::strstr(backend_name, "Vulkan") != nullptr || + std::strstr(backend_name, "CUDA") != nullptr || + std::strstr(backend_name, "ROCm") != nullptr); + const bool direct_depthwise = transcribe::conformer::resolve_conv_direct( + "TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", backend_direct); eb.n_blocks_local = (T_enc + hp.enc_context_size - 1) / hp.enc_context_size; const int T_pad = eb.n_blocks_local * hp.enc_context_size; eb.last_block_rem = T_enc - (eb.n_blocks_local - 1) * hp.enc_context_size; @@ -381,7 +427,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // bias = bn_b - bn_mean * scale // into [inner_dim] tensors stashed under conv_bn_fused_*. ggml_tensor * conv_out = granite_conv_module(ctx, x, b, b.conv_bn_fused_scale, b.conv_bn_fused_bias, conv_k, - static_cast(inner_dim)); + static_cast(inner_dim), direct_depthwise); x = ggml_add(ctx, x, conv_out); // --- FF2 macaron half --- diff --git a/src/arch/granite/encoder.h b/src/arch/granite/encoder.h index 7cefd83a..06a09a17 100644 --- a/src/arch/granite/encoder.h +++ b/src/arch/granite/encoder.h @@ -96,12 +96,14 @@ struct EncoderBuild { // (== n_mel_frames / 2 after the whisper-mode trim). `use_flash` is // reserved for future use — the implementation uses manual mul_mat + soft_max // because the Shaw bias requires a per-(head, block) additive term and -// the flash_attn_ext path doesn't yet broadcast that cleanly. +// the flash_attn_ext path doesn't yet broadcast that cleanly. `backend_name` +// selects the bounded depthwise-convolution implementation. EncoderBuild build_encoder_graph(ggml_context * ctx, const GraniteWeights & weights, const GraniteHParams & hp, int T_enc, - bool use_flash); + bool use_flash, + const char * backend_name); // Host-side precomputation of the Shaw positional-bias rows. // diff --git a/src/arch/granite/model.cpp b/src/arch/granite/model.cpp index 95d182e6..eb34439b 100644 --- a/src/arch/granite/model.cpp +++ b/src/arch/granite/model.cpp @@ -826,7 +826,8 @@ transcribe_status run(transcribe_session * ctx_base, } // Build encoder graph. - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash, + ggml_backend_name(cm->plan.primary)); if (eb.graph == nullptr || eb.out == nullptr) { return TRANSCRIBE_ERR_GGUF; } @@ -1344,7 +1345,8 @@ transcribe_status encode_one(GraniteSession * cc, "out of memory."); return TRANSCRIBE_ERR_OOM; } - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, t_enc, cc->encoder_use_flash); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, t_enc, cc->encoder_use_flash, + ggml_backend_name(cm->plan.primary)); if (eb.graph == nullptr || eb.out == nullptr) { return TRANSCRIBE_ERR_GGUF; } From 931dfb0d87476bdc73b1add51eaba4dda54b9ce2 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Wed, 16 Sep 2026 22:15:59 +0800 Subject: [PATCH 10/17] voxtral mem reduction --- docs/models/voxtral-realtime.md | 30 ++-- docs/porting/families/voxtral_realtime.md | 4 +- include/transcribe/voxtral_realtime.h | 6 +- src/arch/voxtral_realtime/capabilities.cpp | 6 +- src/arch/voxtral_realtime/model.cpp | 148 ++++++++++++++----- src/arch/voxtral_realtime/voxtral_realtime.h | 1 + 6 files changed, 133 insertions(+), 62 deletions(-) diff --git a/docs/models/voxtral-realtime.md b/docs/models/voxtral-realtime.md index d1bdce8a..ca0fc21c 100644 --- a/docs/models/voxtral-realtime.md +++ b/docs/models/voxtral-realtime.md @@ -26,7 +26,9 @@ Real-time and offline speech-to-text from a 16 kHz mono WAV. - **Configurable streaming delay** — `--stream-voxtral-delay ` (default 6 = 480 ms; range 80 ms–2.4 s) sets the transcription delay. - **Accuracy-first offline default** — one-shot and batch inference use the - publisher's most accurate evaluated delay, 30 tokens (2.4 s). + publisher's most accurate evaluated delay, 30 tokens (2.4 s). Long one-shot + inputs are internally advanced through the bounded incremental encoder and + decoder instead of constructing a full-clip attention graph. - Auto language detection (the streaming processor is auto-detect only). ## Input limits @@ -82,11 +84,10 @@ CLI flags: - `--stream-chunk-ms ` — incremental streaming at N-ms chunk granularity. - `--stream-voxtral-delay ` — transcription delay in audio slots (default 6 = 480 ms). -- `--spec-k-drafts ` — offline-path 1-gram-lookup speculative decoding - draft length. `-1` (default) uses the family default (`2`). `0` disables - spec (plain autoregression). `1..8` selects an explicit K. Speculation - applies to `transcribe_run` / `transcribe-cli` only — the streaming path - is unaffected. +- `--spec-k-drafts ` — short-clip offline-path 1-gram-lookup speculative + decoding draft length. `-1` uses the family default (`1`), `0` disables + speculation, and `1..8` selects an explicit K. Long one-shot inputs switch + to the bounded incremental path, which currently uses plain greedy decoding. ## Performance @@ -132,14 +133,15 @@ uv run scripts/bench/run.py \ ## Speculative decoding -The offline decoder runs 1-gram-lookup speculative decoding by default. Each -verify pass processes K+1 positions in parallel: position 0 is the model's -true next-token decision; positions 1..K verify K draft tokens read from the -1-gram suffix lookup over the already-decoded prefix. Drafts are accepted as -long as the model's argmax matches the drafted token; the first mismatch ends -the accepted prefix. Because ~60–70% of audio slots emit `STREAMING_PAD` (id -32), the 1-gram lookup hits high acceptance during silence and during repeated -phrases. +The short-clip offline decoder runs 1-gram-lookup speculative decoding by +default. Each verify pass processes K+1 positions in parallel: position 0 is +the model's true next-token decision; positions 1..K verify K draft tokens read +from the 1-gram suffix lookup over the already-decoded prefix. Drafts are +accepted as long as the model's argmax matches the drafted token; the first +mismatch ends the accepted prefix. Because ~60–70% of audio slots emit +`STREAMING_PAD` (id 32), the 1-gram lookup hits high acceptance during silence +and during repeated phrases. Long one-shot inputs use the bounded incremental +scheduler and plain greedy decoding so encoder and decoder memory stay bounded. The transcript is byte-identical to the K=0 (no-spec) path; only wall-clock time changes. diff --git a/docs/porting/families/voxtral_realtime.md b/docs/porting/families/voxtral_realtime.md index 2914d75e..1bd83b5d 100644 --- a/docs/porting/families/voxtral_realtime.md +++ b/docs/porting/families/voxtral_realtime.md @@ -53,7 +53,9 @@ Pattern: **audio-llm**, streaming, ADDITIVE audio fusion. sliding-KV(8192) re-run incrementally; downsample_factor=4 enc frames per decode step (12.5 Hz); output length clamped to `ceil(mel_frames / audio_length_per_tok=8)`; configurable `num_delay_tokens` - (default 6 = 480 ms). + (default 6 = 480 ms). Long one-shot inference uses this same bounded + scheduler with the offline delay (30), while short clips retain the faster + whole-graph path. ## Family-specific requirements (do not flow through convert/validate) diff --git a/include/transcribe/voxtral_realtime.h b/include/transcribe/voxtral_realtime.h index b2f8fc64..9b32bb3f 100644 --- a/include/transcribe/voxtral_realtime.h +++ b/include/transcribe/voxtral_realtime.h @@ -47,9 +47,9 @@ extern "C" { * min_decode_interval_ms * * Minimum audio-time interval between tentative partial decodes while a - * stream is ACTIVE. Voxtral Realtime's partial decode reprocesses the - * accumulated buffer, so this knob bounds partial-decode compute at the - * cost of less frequent tentative transcripts. stream_finalize always + * stream is ACTIVE. Voxtral Realtime advances its incremental encoder and + * decoder when this interval elapses, so the knob trades partial-result + * latency for larger compute batches. stream_finalize always * performs the final decode regardless of this throttle, and that final * decode is byte-identical to offline inference when both use delay 30. * diff --git a/src/arch/voxtral_realtime/capabilities.cpp b/src/arch/voxtral_realtime/capabilities.cpp index 496bd490..7aa14825 100644 --- a/src/arch/voxtral_realtime/capabilities.cpp +++ b/src/arch/voxtral_realtime/capabilities.cpp @@ -20,9 +20,9 @@ void apply_family_invariants(transcribe_model & model) { // include/transcribe/voxtral_realtime.h. caps.supports_streaming = true; - // Offline path uses 1-gram-lookup speculative decode (verify graph at - // T = spec_k_drafts + 1). Streaming and batched paths are not yet - // spec-enabled. + // The short-clip offline path supports 1-gram-lookup speculative decode + // (verify graph at T = spec_k_drafts + 1). Long one-shot, streaming, and + // batched paths use plain greedy decoding. caps.supports_spec_decode = true; transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); diff --git a/src/arch/voxtral_realtime/model.cpp b/src/arch/voxtral_realtime/model.cpp index 6adfb758..f3d112a2 100644 --- a/src/arch/voxtral_realtime/model.cpp +++ b/src/arch/voxtral_realtime/model.cpp @@ -88,6 +88,12 @@ constexpr const char k_default_variant[] = "voxtral-mini-4b-realtime-2602"; // Offline inference has no latency tradeoff, so use the best evaluated delay. constexpr int k_offline_num_delay_tokens = 30; +// Incremental encoder ring geometry. Keep the trained 750-frame window and +// leave room for bounded batches of new frames before compacting the cache. +constexpr int k_enc_ring_ctx = 1536; +constexpr int k_enc_max_batch = 512; +constexpr int k_offline_incremental_min_frames = 4096; + // Resolve BOS / STREAMING_PAD / EOS against the loaded tokenizer. transcribe_status resolve_specials(const transcribe::Tokenizer & tok, const HParams & hp, PromptSpecials & out) { out.bos = tok.bos_id(); @@ -462,9 +468,8 @@ transcribe_status compute_ada_scales(Session * cc, Model * cm, int num_delay) { return TRANSCRIBE_OK; } -// Core forward: mel -> encoder/projector -> autoregressive decode -> detok. -// Shared by the offline run() and the streaming hooks. `num_delay` drives the -// audio right-pad, the prompt length, and the adaptive-norm scales (they MUST +// Whole-graph forward for short offline inputs and numerical dumps. `num_delay` +// drives the audio right-pad, prompt length, and adaptive-norm scales (they MUST // agree). On success out_text holds the trimmed transcript. Does not touch the // result snapshot (segments / full_text / has_result) — the caller owns that. transcribe_status forward_buffer(Session * cc, @@ -1137,8 +1142,10 @@ transcribe_status forward_buffer(Session * cc, return TRANSCRIBE_OK; } -// Offline one-shot entry point. Thin wrapper over forward_buffer that owns the -// result snapshot (full_text + a single text-only segment). +transcribe_status run_incremental(Session * cc, Model * cm, const float * pcm, int n_samples); + +// Offline one-shot entry point. Short clips retain the whole-graph path (and +// its speculative decoder); longer clips use the bounded streaming scheduler. transcribe_status run(transcribe_session * session, const float * pcm, int n_samples, @@ -1158,6 +1165,16 @@ transcribe_status run(transcribe_session * session, transcribe::debug::init(); const bool dumps_on = transcribe::debug::enabled(); + // For long clips, avoid constructing the quadratic mask and full-clip + // activation graph. Short clips keep the more efficient whole-graph path. + const int64_t raw_per_tok = static_cast(cm->hparams.audio_length_per_tok) * cm->hparams.fe_hop_length; + const int64_t raw_tokens = (static_cast(n_samples) + raw_per_tok - 1) / raw_per_tok; + const int64_t n_audio = cm->specials.n_left_pad + raw_tokens + k_offline_num_delay_tokens + 1 + 10; + const int64_t n_enc = n_audio * cm->hparams.proj_downsample; + if (!dumps_on && n_enc > k_offline_incremental_min_frames) { + return run_incremental(cc, cm, pcm, n_samples); + } + // params->spec_k_drafts: -1 = family default (=1), 0 = disabled, // 1..VOXTRAL_REALTIME_SPEC_K_MAX = explicit. Clamp into range so a // misconfigured caller doesn't ask for an unbounded verify graph. Default @@ -1327,14 +1344,8 @@ bool stream_run_graph(Session * cc, return ok; } -// Encoder KV ring geometry. Keep the last `sliding_window`(750) frames: hold a -// contiguous cache of k_enc_ring_ctx slots, append at the write head, and -// periodically COMPACT (copy the last 750 frames to the front) so the ring never -// overflows. k_enc_max_batch bounds frames per chunk (must be <= ring - -// sliding_window and a multiple of proj_downsample). -constexpr int k_enc_ring_ctx = 1536; -constexpr int k_enc_max_batch = 512; - +// Encoder KV ring: periodically compact the last sliding-window frames to the +// front so the next bounded batch fits. // (decoder sliding-window ring size is read from the GGUF: hp.dec_sliding_window) // Compact the encoder KV ring: move the last `keep` frames (slots @@ -1707,7 +1718,10 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * } if (cc->stream_prompt_done && !cc->stream_eos && cc->stream_dec_pos < cc->stream_n_tok_ready) { - const int max_n_kv = cc->kv_cache.n_ctx; + // Before the decoder ring fills, expose only the populated prefix. The + // remaining slots are all masked and needlessly widen every attention + // step on short and medium streams. + const int max_n_kv = std::min(cc->stream_dec_ring_ctx, cc->stream_n_tok_ready); if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; @@ -1743,8 +1757,8 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * } apply_threads(cc->sched, cc->n_threads); - const int swin = hp.dec_sliding_window; // 8192 - const int kv_ring = max_n_kv; // ring size == n_ctx (== swin here) + const int swin = hp.dec_sliding_window; // 8192 + const int kv_ring = cc->stream_dec_ring_ctx; // physical ring width // Hard absolute-position cap = dec_max_position; a streaming caller hits // memory/latency long before it (~2.9 h). const int max_pos = voxtral_realtime_abs_position_cap(hp); @@ -1775,8 +1789,8 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * const int64_t kv_idx = cur % kv_ring; ggml_backend_tensor_set(sb.kv_idx_in, &kv_idx, 0, sizeof(int64_t)); // Reveal keys [max(0,cur-swin+1) .. cur] at their ring slots (t % kv_ring). - // For cur < kv_ring this is the identity slot map (== the pre-ring path); - // once full, the `swin` in-window tokens occupy all `kv_ring` slots 1:1. + // Before the ring fills, max_n_kv covers the populated identity-mapped + // prefix. Once full, all in-window tokens occupy the ring slots 1:1. std::fill(step_mask.begin(), step_mask.end(), mn); for (int t = std::max(0, cur - swin + 1); t <= cur; ++t) { step_mask[t % kv_ring] = mz; @@ -1868,7 +1882,8 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * } // ---- 7. Publish transcript. ---- - std::string text = detok_generated(cm, cc->stream_generated); + std::string text = detok_generated(cm, cc->stream_generated); + cc->raw_text = cm->tok.decode(cc->stream_generated.data(), static_cast(cc->stream_generated.size())); const int sr = std::max(1, hp.fe_sample_rate); const int64_t audio_ms = static_cast(n_aud_total) * 1000 / sr; publish_stream_text(cc, cm, text, audio_ms); @@ -1878,18 +1893,11 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * return TRANSCRIBE_OK; } -transcribe_status stream_begin(transcribe_session * session, - const transcribe_run_params * /*run_params*/, - const transcribe_stream_params * stream_params) { - auto * cc = static_cast(session); - auto * cm = static_cast(cc->model); - if (cm == nullptr || cm->plan.scheduler_list.empty()) { +transcribe_status begin_stream_state(Session * cc, Model * cm, int nd, int md, int dec_ring, ggml_type enc_kv_type) { + if (cc == nullptr || cm == nullptr || cm->plan.scheduler_list.empty() || dec_ring <= 0 || + dec_ring > cm->hparams.dec_sliding_window) { return TRANSCRIBE_ERR_INVALID_ARG; } - int nd = 0, md = 0; - if (auto st = resolve_stream_ext(stream_params, cm, &nd, &md); st != TRANSCRIBE_OK) { - return st; - } cc->stream_pcm.clear(); cc->stream_num_delay = nd; @@ -1900,6 +1908,7 @@ transcribe_status stream_begin(transcribe_session * session, cc->stream_n_enc_committed = 0; cc->stream_enc_slot = 0; cc->stream_enc_abs_base = 0; + cc->stream_dec_ring_ctx = dec_ring; cc->stream_conv0_cache.assign(static_cast(cm->hparams.enc_num_mel_bins) * 2, 0.0f); cc->stream_conv1_cache.assign(static_cast(cm->hparams.enc_d_model), 0.0f); cc->stream_n_mel_committed = 0; @@ -1915,13 +1924,16 @@ transcribe_status stream_begin(transcribe_session * session, cc->stream_eos = false; cc->stream_generated.clear(); cc->stream_n_audio_clamp = -1; + cc->stream_gen0_logits.clear(); + cc->stream_gen8_logits.clear(); - // Encoder StaticCache ring (F32; MHA so n_kv_heads == n_heads). Fixed + // Encoder StaticCache ring (MHA so n_kv_heads == n_heads). Fixed // k_enc_ring_ctx slots; compaction keeps the last sliding_window(750) frames // so any stream length runs in constant memory (the reference mechanism). + const ggml_type dec_kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; cc->enc_kv.free(); if (!transcribe::causal_lm::kv_init(cc->enc_kv, cm->plan.primary, /*n_ctx=*/k_enc_ring_ctx, cm->hparams.enc_n_heads, - cm->hparams.enc_head_dim, cm->hparams.enc_n_layers, GGML_TYPE_F32)) { + cm->hparams.enc_head_dim, cm->hparams.enc_n_layers, enc_kv_type)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral_realtime stream_begin: encoder KV cache allocation failed — " "out of memory."); @@ -1931,18 +1943,14 @@ transcribe_status stream_begin(transcribe_session * session, ggml_backend_buffer_clear(cc->enc_kv.buffer, 0); } - // Decoder KV: a sliding-window RING sized to the model's own `sliding_window` - // (from the GGUF, 8192 here). The step loop writes token `cur` at slot - // `cur % n_ctx`, so the cache holds the last `swin` tokens for any stream - // length in constant memory. `sliding_window` is a trained-in constant, not - // an inference knob. Sized once per session; never shrinks a larger ctx a - // prior offline run may have left. - const int dec_ring = cm->hparams.dec_sliding_window; + // Decoder KV is a sliding ring. Public streams use the trained window; + // one-shot runs may use a smaller ring when the known horizon fits in it. + // The backing allocation only grows, while stream_dec_ring_ctx selects the + // active prefix so results do not depend on prior runs in this session. if (cc->kv_cache.n_ctx < dec_ring) { - const ggml_type kt = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; cc->kv_cache.free(); if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, dec_ring, cm->hparams.dec_n_kv_heads, - cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, kt)) { + cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, dec_kv_type)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "voxtral_realtime stream_begin: decoder KV cache allocation " "failed — out of memory."); @@ -1964,6 +1972,21 @@ transcribe_status stream_begin(transcribe_session * session, return TRANSCRIBE_OK; } +transcribe_status stream_begin(transcribe_session * session, + const transcribe_run_params * /*run_params*/, + const transcribe_stream_params * stream_params) { + auto * cc = static_cast(session); + auto * cm = static_cast(cc->model); + if (cm == nullptr || cm->plan.scheduler_list.empty()) { + return TRANSCRIBE_ERR_INVALID_ARG; + } + int nd = 0, md = 0; + if (auto st = resolve_stream_ext(stream_params, cm, &nd, &md); st != TRANSCRIBE_OK) { + return st; + } + return begin_stream_state(cc, cm, nd, md, cm->hparams.dec_sliding_window, GGML_TYPE_F32); +} + transcribe_status stream_feed(transcribe_session * session, const float * pcm, int n_samples, @@ -2019,8 +2042,12 @@ transcribe_status stream_finalize(transcribe_session * session, transcribe_strea return st; } + cc->t_mel_us = cc->stream_t_mel_us; + cc->t_encode_us = cc->stream_t_conv_us + cc->stream_t_enc_us; + cc->t_decode_us = cc->stream_t_dec_us; + const int sr = std::max(1, cm->hparams.fe_sample_rate); - const int64_t audio_ms = static_cast(cc->stream_pcm.size()) * 1000 / sr; + const int64_t audio_ms = (cc->stream_pcm_drop + static_cast(cc->stream_pcm.size())) * 1000 / sr; if (update != nullptr) { update->result_changed = true; update->revision = cc->stream_revision; @@ -2031,6 +2058,44 @@ transcribe_status stream_finalize(transcribe_session * session, transcribe_strea return TRANSCRIBE_OK; } +transcribe_status run_incremental(Session * cc, Model * cm, const float * pcm, int n_samples) { + const int64_t raw_per_tok = static_cast(cm->hparams.audio_length_per_tok) * cm->hparams.fe_hop_length; + const int64_t raw_tokens = (static_cast(n_samples) + raw_per_tok - 1) / raw_per_tok; + const int64_t n_audio = cm->specials.n_left_pad + raw_tokens + k_offline_num_delay_tokens + 1 + 10; + const int abs_cap = voxtral_realtime_abs_position_cap(cm->hparams); + if (n_audio + 1 > abs_cap) { + transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "voxtral_realtime run: clip needs %lld positions > model max %d " + "(dec_max_position, ~2.9 h). See transcribe_capabilities.max_audio_ms.", + static_cast(n_audio + 1), abs_cap); + return TRANSCRIBE_ERR_INPUT_TOO_LONG; + } + + int dec_ring = 2048; + while (dec_ring < n_audio + 1 && dec_ring < cm->hparams.dec_sliding_window) { + dec_ring *= 2; + } + dec_ring = std::min(dec_ring, cm->hparams.dec_sliding_window); + const ggml_type enc_kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; + if (auto st = begin_stream_state(cc, cm, k_offline_num_delay_tokens, /*min_decode_ms=*/0, dec_ring, enc_kv_type); + st != TRANSCRIBE_OK) { + return st; + } + + // Feed one maximum encoder batch of real audio at a time. Initial/final + // padding can spill into one additional bounded encoder graph. + const int chunk_samples = k_enc_max_batch * 2 * cm->hparams.fe_hop_length; + int pos = 0; + while (pos < n_samples) { + const int take = std::min(chunk_samples, n_samples - pos); + if (auto st = stream_feed(cc, pcm + pos, take, nullptr); st != TRANSCRIBE_OK) { + return st; + } + pos += take; + } + return stream_finalize(cc, nullptr); +} + void stream_reset(transcribe_session * session) { auto * cc = static_cast(session); cc->stream_pcm.clear(); @@ -2040,6 +2105,7 @@ void stream_reset(transcribe_session * session) { cc->stream_n_enc_committed = 0; cc->stream_enc_slot = 0; cc->stream_enc_abs_base = 0; + cc->stream_dec_ring_ctx = 0; cc->stream_conv0_cache.clear(); cc->stream_conv1_cache.clear(); cc->stream_n_mel_committed = 0; diff --git a/src/arch/voxtral_realtime/voxtral_realtime.h b/src/arch/voxtral_realtime/voxtral_realtime.h index 88a1ef49..d0495c5e 100644 --- a/src/arch/voxtral_realtime/voxtral_realtime.h +++ b/src/arch/voxtral_realtime/voxtral_realtime.h @@ -101,6 +101,7 @@ struct Session final : public transcribe_session { int stream_enc_slot = 0; // ring write head (slot units) int stream_enc_abs_base = 0; // absolute frame index of slot 0 int stream_n_enc_committed = 0; // enc frames committed (absolute) + int stream_dec_ring_ctx = 0; // active decoder ring width // Conv-stem padding cache: feed only new mel frames, carrying the conv // left-context. conv0 (k3 s1) ← last 2 mel frames; conv1 (k3 s2) ← last 1 // conv0-output frame. Zeros on the first chunk == whole-buffer left-pad. From c9831db02af51920cc824537d804f56ad5856f59 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 07:29:36 +0800 Subject: [PATCH 11/17] sanm mem reduction --- src/arch/funasr_nano/encoder.cpp | 13 +++- src/arch/funasr_nano/encoder.h | 3 +- src/arch/funasr_nano/model.cpp | 28 ++------- src/arch/sensevoice/encoder.cpp | 63 ++++++++++++++----- src/arch/sensevoice/encoder.h | 6 +- src/arch/sensevoice/model.cpp | 104 +++++++++++++++++-------------- src/arch/sensevoice/sensevoice.h | 3 +- src/sanm/sanm.cpp | 71 +++++++++++++++------ src/sanm/sanm.h | 12 ++-- src/transcribe-batch-util.cpp | 25 ++++++++ src/transcribe-batch-util.h | 10 +++ 11 files changed, 224 insertions(+), 114 deletions(-) diff --git a/src/arch/funasr_nano/encoder.cpp b/src/arch/funasr_nano/encoder.cpp index dfe11594..e27a13a8 100644 --- a/src/arch/funasr_nano/encoder.cpp +++ b/src/arch/funasr_nano/encoder.cpp @@ -59,7 +59,8 @@ void mark_dump(ggml_tensor *& slot, ggml_tensor * t, const char * name) { EncoderBuild build_encoder_graph(ggml_context * ctx, const FunAsrNanoWeights & w, const FunAsrNanoHParams & hp, - int n_lfr_frames) { + int n_lfr_frames, + const char * backend_name) { EncoderBuild eb{}; if (ctx == nullptr || n_lfr_frames <= 0) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -73,11 +74,17 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const int d_model = hp.enc_d_model; const int T = n_lfr_frames; - const sanm::SanmBlockParams block_params{ + sanm::SanmBlockParams block_params{ /*n_heads=*/hp.enc_n_heads, /*d_model=*/d_model, /*kernel=*/hp.enc_kernel, }; + const bool backend_direct = backend_name != nullptr && (std::strstr(backend_name, "Vulkan") != nullptr || + std::strstr(backend_name, "CUDA") != nullptr || + std::strstr(backend_name, "ROCm") != nullptr); + block_params.direct_depthwise = + conf::resolve_conv_direct("TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", backend_direct); + block_params.bounded_depthwise = true; eb.frontend_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_input, T); named(eb.frontend_in, "frontend.in"); @@ -147,7 +154,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.out = x; ggml_set_output(eb.out); - eb.graph = ggml_new_graph_custom(ctx, /*size=*/8192, /*grads=*/false); + eb.graph = ggml_new_graph_custom(ctx, /*size=*/32768, /*grads=*/false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "funasr_nano encoder: ggml_new_graph_custom failed"); return eb; diff --git a/src/arch/funasr_nano/encoder.h b/src/arch/funasr_nano/encoder.h index 9772d405..a81e71a6 100644 --- a/src/arch/funasr_nano/encoder.h +++ b/src/arch/funasr_nano/encoder.h @@ -48,6 +48,7 @@ struct EncoderBuild { EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const FunAsrNanoWeights & weights, const FunAsrNanoHParams & hp, - int n_lfr_frames); + int n_lfr_frames, + const char * backend_name = nullptr); } // namespace transcribe::funasr_nano diff --git a/src/arch/funasr_nano/model.cpp b/src/arch/funasr_nano/model.cpp index 436ec1ce..39b0f7e6 100644 --- a/src/arch/funasr_nano/model.cpp +++ b/src/arch/funasr_nano/model.cpp @@ -425,24 +425,6 @@ transcribe_status init_context(transcribe_model * model, cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - auto * cm = static_cast(model); - { - ggml_type kv_type = GGML_TYPE_F16; - if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { - kv_type = GGML_TYPE_F32; - } - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, - /*n_ctx=*/2048, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, - cm->hparams.dec_n_layers, kv_type)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "funasr_nano init_context: KV cache allocation failed " - "(n_ctx=2048, %d kv-heads x %d head-dim x %d layers) — " - "out of memory.", - cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, cm->hparams.dec_n_layers); - return TRANSCRIBE_ERR_OOM; - } - } - *out_ctx = cc.release(); return TRANSCRIBE_OK; } @@ -512,14 +494,14 @@ transcribe_status run(transcribe_session * session, } // ---- Build encoder graph ---- - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr, cm->backend.c_str()); if (eb.graph == nullptr || eb.out == nullptr) { return TRANSCRIBE_ERR_GGUF; } if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, - static_cast(cm->plan.scheduler_list.size()), 16384, /*parallel=*/false, + static_cast(cm->plan.scheduler_list.size()), 32768, /*parallel=*/false, /*op_offload=*/true); if (cc->sched == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "funasr_nano run: ggml_backend_sched_new failed"); @@ -681,7 +663,7 @@ transcribe_status run(transcribe_session * session, // power of two (the step graph's attention width wants pow2 for the fast // flash-attn path). The cache grows across runs as audio length demands; // a pre-allocated smaller cache is freed and re-allocated. - int want_n_ctx = 1024; + int want_n_ctx = 256; while (want_n_ctx < T_prompt + k_max_new) { want_n_ctx *= 2; } @@ -978,13 +960,13 @@ transcribe_status audio_embed_one(FunAsrNanoSession * cc, if (const transcribe_status st = reset_ctx(cc, 32); st != TRANSCRIBE_OK) { return st; } - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr, cm->backend.c_str()); if (eb.graph == nullptr || eb.out == nullptr) { return TRANSCRIBE_ERR_GGUF; } if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, - static_cast(cm->plan.scheduler_list.size()), 16384, false, true); + static_cast(cm->plan.scheduler_list.size()), 32768, false, true); if (cc->sched == nullptr) { return TRANSCRIBE_ERR_GGUF; } diff --git a/src/arch/sensevoice/encoder.cpp b/src/arch/sensevoice/encoder.cpp index 0715ff49..9b35734c 100644 --- a/src/arch/sensevoice/encoder.cpp +++ b/src/arch/sensevoice/encoder.cpp @@ -22,9 +22,11 @@ #include "transcribe-log.h" #include "weights.h" +#include #include #include #include +#include namespace transcribe::sensevoice { @@ -67,7 +69,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const SenseVoiceHParams & hp, int n_lfr_frames, int n_batch, - bool batch_var_len) { + bool batch_var_len, + const char * backend_name) { EncoderBuild eb{}; if (ctx == nullptr || n_lfr_frames <= 0 || n_batch <= 0) { @@ -90,6 +93,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, /*d_model=*/d_model, /*kernel=*/hp.enc_kernel, }; + const bool backend_direct = backend_name != nullptr && (std::strstr(backend_name, "Vulkan") != nullptr || + std::strstr(backend_name, "CUDA") != nullptr || + std::strstr(backend_name, "ROCm") != nullptr); + block_params.direct_depthwise = transcribe::conformer::resolve_conv_direct( + "TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", backend_direct); + block_params.bounded_depthwise = true; // ----- inputs ---------------------------------------------------- // Batch axis at ne[2]. B == 1 collapses to the pre-batch [d_input, T_in] @@ -230,21 +239,47 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, mark_dump(eb.dumps.tp_norm_out, x, "enc.tp_norm.out"); // ----- CTC head -------------------------------------------------- - ggml_tensor * logits = ggml_mul_mat(ctx, w.ctc_head_w, x); - logits = ggml_add(ctx, logits, w.ctc_head_b); - mark_dump(eb.dumps.ctc_logits, logits, "ctc.logits.raw"); - - // log_softmax over the vocab axis (ne[0]). - ggml_tensor * log_probs = ggml_log(ctx, ggml_soft_max(ctx, logits)); - mark_dump(eb.dumps.ctc_log_probs, log_probs, "ctc.log_probs"); - - eb.out = log_probs; + if (transcribe::debug::enabled()) { + // Numerical validation retains the complete reference tensors. + ggml_tensor * logits = ggml_mul_mat(ctx, w.ctc_head_w, x); + logits = ggml_add(ctx, logits, w.ctc_head_b); + mark_dump(eb.dumps.ctc_logits, logits, "ctc.logits.raw"); + + ggml_tensor * log_probs = ggml_log(ctx, ggml_soft_max(ctx, logits)); + mark_dump(eb.dumps.ctc_log_probs, log_probs, "ctc.log_probs"); + eb.out = log_probs; + } else { + // Greedy CTC needs only one id per frame. Project bounded groups of + // columns so neither the complete [vocab,T,B] logits nor log-softmax + // matrix becomes persistent backend workspace. + constexpr int64_t kCtcColumnsPerChunk = 256; + const int64_t n_columns = x->ne[1] * x->ne[2]; + ggml_tensor * x_flat = ggml_reshape_2d(ctx, ggml_cont(ctx, x), x->ne[0], n_columns); + std::vector chunks; + chunks.reserve(static_cast((n_columns + kCtcColumnsPerChunk - 1) / kCtcColumnsPerChunk)); + for (int64_t col = 0; col < n_columns; col += kCtcColumnsPerChunk) { + const int64_t n = std::min(kCtcColumnsPerChunk, n_columns - col); + ggml_tensor * x_chunk = ggml_view_2d(ctx, x_flat, x_flat->ne[0], n, x_flat->nb[1], col * x_flat->nb[1]); + ggml_tensor * logits = ggml_mul_mat(ctx, w.ctc_head_w, x_chunk); + logits = ggml_add(ctx, logits, w.ctc_head_b); + chunks.push_back(ggml_argmax(ctx, logits)); + } + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/0) : + chunks[i]); + } + chunks.swap(next); + } + eb.out = chunks.front(); + } ggml_set_output(eb.out); - // 70 SAN-M blocks * ~30 ops/block + frontend/PE/CTC ≈ 2300 nodes; the - // variable-length batch path's manual SDPA adds a handful more per block. - // 8192 leaves ample headroom. - eb.graph = ggml_new_graph_custom(ctx, /*size=*/8192, /*grads=*/false); + // Long inputs split every FSMN convolution into bounded im2col chunks; + // reserve enough nodes for all 70 blocks without scaling workspace. + eb.graph = ggml_new_graph_custom(ctx, /*size=*/32768, /*grads=*/false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sensevoice encoder: ggml_new_graph_custom failed"); return eb; diff --git a/src/arch/sensevoice/encoder.h b/src/arch/sensevoice/encoder.h index 26a32413..54501fcb 100644 --- a/src/arch/sensevoice/encoder.h +++ b/src/arch/sensevoice/encoder.h @@ -85,7 +85,8 @@ struct EncoderBuild { ggml_tensor * attn_pad_mask_in = nullptr; ggml_tensor * conv_pad_mask_in = nullptr; - // CTC log-probabilities. ne=[vocab, T, n_batch, 1]. + // Debug mode: CTC log-probabilities [vocab, T, n_batch]. Normal mode: + // per-frame argmax ids [T * n_batch] i32 from a chunked CTC head. ggml_tensor * out = nullptr; EncoderDumps dumps{}; @@ -104,6 +105,7 @@ EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const SenseVoiceHParams & hp, int n_lfr_frames, int n_batch = 1, - bool batch_var_len = false); + bool batch_var_len = false, + const char * backend_name = nullptr); } // namespace transcribe::sensevoice diff --git a/src/arch/sensevoice/model.cpp b/src/arch/sensevoice/model.cpp index d0b49f17..cc2ea5dc 100644 --- a/src/arch/sensevoice/model.cpp +++ b/src/arch/sensevoice/model.cpp @@ -238,20 +238,26 @@ void apply_thread_policy(SenseVoiceSession * cc) { transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } -// Greedy CTC decode + public result-hierarchy build for ONE utterance's CTC -// log-probabilities. Shared by the single-shot path (run) and the batched -// path (run_batch). `log_probs` is the row-major [T_full, vocab] buffer for -// this utterance (row t at log_probs + t*vocab); `lang` is the caller's -// requested language (or null/auto) and `n_samples` sizes the segment -// duration. Writes the session scratch result slot (token_ids / tokens / -// segments / full_text / detected_language / result_kind / has_result) and -// records cc->t_decode_us. +static void argmax_rows(const float * rows, int n_rows, int n_columns, int32_t * out) { + for (int t = 0; t < n_rows; ++t) { + const float * row = rows + static_cast(t) * n_columns; + int32_t best = 0; + for (int v = 1; v < n_columns; ++v) { + if (row[v] > row[best]) { + best = v; + } + } + out[t] = best; + } +} + +// Greedy CTC collapse + public result-hierarchy build for one utterance. +// `frame_ids` is the device- or host-computed argmax at each encoder frame. static transcribe_status decode_and_populate(SenseVoiceSession * cc, SenseVoiceModel * cm, const transcribe_run_params * params, - const float * log_probs, + const int32_t * frame_ids, int T_full, - int vocab, const char * lang, int n_samples) { const auto & hp = cm->hparams; @@ -263,20 +269,12 @@ static transcribe_status decode_and_populate(SenseVoiceSession * cc, cc->token_ids.reserve(static_cast(T_full)); int prev_id = -1; for (int t = 0; t < T_full; ++t) { - const float * row = log_probs + static_cast(t) * vocab; - int best_id = 0; - float best = row[0]; - for (int v = 1; v < vocab; ++v) { - if (row[v] > best) { - best = row[v]; - best_id = v; - } - } - if (best_id != prev_id) { - if (best_id != blank_id) { - cc->token_ids.push_back(best_id); + const int32_t id = frame_ids[t]; + if (id != prev_id) { + if (id != blank_id) { + cc->token_ids.push_back(id); } - prev_id = best_id; + prev_id = id; } } @@ -451,7 +449,8 @@ transcribe_status run(transcribe_session * session, } // ---------- Build the encoder graph ------------------------------- - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_lfr, /*n_batch=*/1, + /*batch_var_len=*/false, cm->backend.c_str()); if (eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; } @@ -460,7 +459,7 @@ transcribe_status run(transcribe_session * session, if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, static_cast(cm->plan.scheduler_list.size()), - /*graph_size=*/8192, /*parallel=*/false, /*op_offload=*/true); + /*graph_size=*/32768, /*parallel=*/false, /*op_offload=*/true); if (cc->sched == nullptr) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sensevoice run: scheduler allocation failed — out of memory. " @@ -575,13 +574,18 @@ transcribe_status run(transcribe_session * session, try_dump("ctc.logits.raw", eb.dumps.ctc_logits, "ctc.logits.raw"); try_dump("ctc.log_probs", eb.dumps.ctc_log_probs, "ctc.log_probs"); - // ---------- Read CTC log-probs to host --------------------------- + // ---------- Read greedy CTC ids ---------------------------------- const int vocab = hp.vocab_size; - cc->logits_buf.resize(static_cast(T_full) * vocab); - ggml_backend_tensor_get(eb.out, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); + cc->argmax_buf.resize(static_cast(T_full)); + if (transcribe::debug::enabled()) { + cc->logits_buf.resize(static_cast(T_full) * vocab); + ggml_backend_tensor_get(eb.out, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); + argmax_rows(cc->logits_buf.data(), T_full, vocab, cc->argmax_buf.data()); + } else { + ggml_backend_tensor_get(eb.out, cc->argmax_buf.data(), 0, cc->argmax_buf.size() * sizeof(int32_t)); + } - // ---------- Greedy CTC decode + public result -------------------- - return decode_and_populate(cc, cm, params, cc->logits_buf.data(), T_full, vocab, lang, n_samples); + return decode_and_populate(cc, cm, params, cc->argmax_buf.data(), T_full, lang, n_samples); } // --------------------------------------------------------------------------- @@ -677,7 +681,7 @@ static transcribe_status run_batch_encode( } EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_max_lfr, /*n_batch=*/n, - /*batch_var_len=*/var_len); + /*batch_var_len=*/var_len, cm->backend.c_str()); if (eb.out == nullptr || eb.graph == nullptr || eb.frontend_in == nullptr) { return TRANSCRIBE_ERR_GGUF; } @@ -685,7 +689,7 @@ static transcribe_status run_batch_encode( if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new(cm->plan.scheduler_list.data(), nullptr, static_cast(cm->plan.scheduler_list.size()), - /*graph_size=*/8192, /*parallel=*/false, /*op_offload=*/true); + /*graph_size=*/32768, /*parallel=*/false, /*op_offload=*/true); if (cc->sched == nullptr) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "sensevoice run: scheduler allocation failed — out of memory. " @@ -764,26 +768,30 @@ static transcribe_status run_batch_encode( } cc->t_encode_us = ggml_time_us() - t_enc_start; - // ---------- Read CTC log-probs + per-utterance decode ------------ - const int vocab = hp.vocab_size; - const size_t utt_elems = static_cast(vocab) * static_cast(T_full_max); - cc->logits_buf.resize(utt_elems * static_cast(n)); - ggml_backend_tensor_get(eb.out, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); - - // Host-slice the shared CTC log-probs and decode each utterance, with the - // single shared encode + total mel cost amortized across the batch. - return transcribe::decode_batch_slices( - cc, n, cc->logits_buf.data(), utt_elems, cc->t_encode_us, total_mel_us, [&](int b, const float * lp) { - // Per-utterance CTC log-probs dump for the batch tensor-parity - // gate. Same vocab-innermost element order as the single-shot - // ctc.log_probs dump, so the harness can diff slice-for-slice. - if (transcribe::debug::enabled()) { + // ---------- Read CTC output + per-utterance decode --------------- + const int vocab = hp.vocab_size; + if (transcribe::debug::enabled()) { + const size_t utt_elems = static_cast(vocab) * static_cast(T_full_max); + cc->logits_buf.resize(utt_elems * static_cast(n)); + ggml_backend_tensor_get(eb.out, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); + cc->argmax_buf.resize(static_cast(T_full_max)); + return transcribe::decode_batch_slices( + cc, n, cc->logits_buf.data(), utt_elems, cc->t_encode_us, total_mel_us, [&](int b, const float * lp) { const long long shape[2] = { real_T_full[b], vocab }; std::string nm = "ctc.log_probs.b" + std::to_string(b); transcribe::debug::dump_host_f32(nm.c_str(), lp, static_cast(real_T_full[b]) * vocab, shape, 2, "ctc.log_probs"); - } - return decode_and_populate(cc, cm, params, lp, real_T_full[b], vocab, lang, n_samples[b]); + argmax_rows(lp, real_T_full[b], vocab, cc->argmax_buf.data()); + return decode_and_populate(cc, cm, params, cc->argmax_buf.data(), real_T_full[b], lang, n_samples[b]); + }); + } + + const size_t utt_ids = static_cast(T_full_max); + cc->argmax_buf.resize(utt_ids * static_cast(n)); + ggml_backend_tensor_get(eb.out, cc->argmax_buf.data(), 0, cc->argmax_buf.size() * sizeof(int32_t)); + return transcribe::decode_batch_id_slices( + cc, n, cc->argmax_buf.data(), utt_ids, cc->t_encode_us, total_mel_us, [&](int b, const int32_t * ids) { + return decode_and_populate(cc, cm, params, ids, real_T_full[b], lang, n_samples[b]); }); } diff --git a/src/arch/sensevoice/sensevoice.h b/src/arch/sensevoice/sensevoice.h index 09564c8e..be33df6b 100644 --- a/src/arch/sensevoice/sensevoice.h +++ b/src/arch/sensevoice/sensevoice.h @@ -52,7 +52,8 @@ struct SenseVoiceSession final : public transcribe_session { // Reusable host scratch. std::vector frontend_buf; // [T_lfr, d_input] std::vector pe_buf; // [T, d_input] - std::vector logits_buf; // [T, vocab] for greedy CTC + std::vector logits_buf; // debug-only [T, vocab] CTC log-probs + std::vector argmax_buf; // normal path [T] greedy CTC ids std::vector token_ids; // post-collapse / post-blank-strip SenseVoiceSession() = default; diff --git a/src/sanm/sanm.cpp b/src/sanm/sanm.cpp index fb231f6c..f82c166f 100644 --- a/src/sanm/sanm.cpp +++ b/src/sanm/sanm.cpp @@ -7,7 +7,9 @@ #include "conformer/conformer.h" #include "ggml.h" +#include #include +#include namespace transcribe::sanm { @@ -36,12 +38,14 @@ ggml_tensor * fsmn_branch(ggml_context * ctx, ggml_tensor * v_pre, // [d_model, T, B] ggml_tensor * fsmn_w, // ne=[K, 1, d_model] int kernel, - ggml_tensor * conv_pad_mask) // [1, T, B] or null -{ + ggml_tensor * conv_pad_mask, // [1, T, B] or null + bool direct_depthwise, + bool bounded_depthwise) { namespace conf = transcribe::conformer; const int64_t B = v_pre->ne[2]; const int64_t d_model = v_pre->ne[0]; + const int64_t T = v_pre->ne[1]; const int padding = (kernel - 1) / 2; // sanm_shift=0 // Variable-length batch: zero padded frames before the conv so a padded @@ -59,28 +63,58 @@ ggml_tensor * fsmn_branch(ggml_context * ctx, ggml_tensor * v_t = ggml_cont(ctx, ggml_transpose(ctx, v_in)); // [T, d_model, B] ggml_tensor * fsmn; - if (B > 1) { - // Batched depthwise conv. conv_1d_dw_f32 (im2col) collapses the - // batch axis, so use the direct depthwise-2D op (W=T, H=1, C=d_model, - // N=B), which threads the utterance batch at ne[3]. Kernel - // [K, 1, d_model] -> [K, 1, 1, d_model]. + if (B > 1 || direct_depthwise) { + // The direct op supports a real batch axis and avoids im2col's K-fold + // expansion. Promote non-F32 kernels for backend portability. ggml_tensor * knl = ggml_reshape_4d(ctx, fsmn_w, kernel, 1, 1, d_model); ggml_tensor * data = ggml_reshape_4d(ctx, v_t, v_t->ne[0], 1, d_model, B); - fsmn = ggml_conv_2d_dw_direct(ctx, knl, data, - /*s0=*/1, /*s1=*/1, - /*p0=*/padding, /*p1=*/0, - /*d0=*/1, /*d1=*/1); - // [T, 1, d_model, B] -> [T, d_model, B]. + fsmn = conf::conv_2d_dw_direct_f32(ctx, knl, data, + /*s0=*/1, /*s1=*/1, + /*p0=*/padding, /*p1=*/0, + /*d0=*/1, /*d1=*/1); fsmn = ggml_reshape_3d(ctx, fsmn, fsmn->ne[0], d_model, B); - } else { - // Single-shot: im2col path. + if (B == 1) { + fsmn = ggml_reshape_2d(ctx, fsmn, fsmn->ne[0], d_model); + } + } else if (!bounded_depthwise || T <= 256) { + // Keep established CPU/Metal numerics for short validation clips. fsmn = conf::conv_1d_dw_f32(ctx, fsmn_w, v_t, /*stride=*/1, /*padding=*/padding, /*dilation=*/1); - // fsmn ne=[T, d_model, 1]. Drop the singleton batch. - fsmn = ggml_reshape_2d(ctx, fsmn, fsmn->ne[0], fsmn->ne[1]); // [T, d_model] + fsmn = ggml_reshape_2d(ctx, fsmn, fsmn->ne[0], fsmn->ne[1]); + } else { + // Bound im2col to 256 output frames while preserving the full kernel + // receptive field at every chunk boundary. + constexpr int64_t kTimeChunk = 256; + std::vector chunks; + chunks.reserve(static_cast((T + kTimeChunk - 1) / kTimeChunk)); + for (int64_t out0 = 0; out0 < T; out0 += kTimeChunk) { + const int64_t n_out = std::min(kTimeChunk, T - out0); + const int64_t src0 = out0 - padding; + const int64_t src1 = src0 + n_out + kernel - 1; + const int64_t view0 = std::max(0, src0); + const int64_t view1 = std::min(T, src1); + const int64_t pad_left = view0 - src0; + const int64_t pad_right = src1 - view1; + ggml_tensor * input = + ggml_view_3d(ctx, v_t, view1 - view0, d_model, 1, v_t->nb[1], v_t->nb[2], view0 * v_t->nb[0]); + input = ggml_pad_ext(ctx, input, pad_left, pad_right, 0, 0, 0, 0, 0, 0); + chunks.push_back(conf::conv_1d_dw_f32(ctx, fsmn_w, input, + /*stride=*/1, /*padding=*/0, + /*dilation=*/1)); + } + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/0) : + chunks[i]); + } + chunks.swap(next); + } + fsmn = ggml_reshape_2d(ctx, chunks.front(), T, d_model); } - fsmn = ggml_cont(ctx, ggml_transpose(ctx, fsmn)); // [d_model, T, B] + fsmn = ggml_cont(ctx, ggml_transpose(ctx, fsmn)); // [d_model, T, B] // Residual within the FSMN: x_fsmn += masked_v. (Padded frames carry // unmasked residual here, but they are masked out of attention and @@ -124,7 +158,8 @@ ggml_tensor * sanm_attention(ggml_context * ctx, ggml_tensor * v_pre = ggml_cont(ctx, v); // FSMN branch (parallel to SDPA). - ggml_tensor * fsmn = fsmn_branch(ctx, v_pre, b.attn_fsmn_w, kernel, p.conv_pad_mask); + ggml_tensor * fsmn = + fsmn_branch(ctx, v_pre, b.attn_fsmn_w, kernel, p.conv_pad_mask, p.direct_depthwise, p.bounded_depthwise); // SDPA. Reshape Q,K,V to [head_dim, n_heads, T, B]. The fused QKV split // is non-contiguous along the channel axis, so cont each before reshape. diff --git a/src/sanm/sanm.h b/src/sanm/sanm.h index 4dbe4d97..6de1c37f 100644 --- a/src/sanm/sanm.h +++ b/src/sanm/sanm.h @@ -69,9 +69,11 @@ struct SanmBlockParams { int d_model = 0; int kernel = 0; // FSMN depthwise kernel width (sanm_shift=0) - ggml_tensor * attn_pad_mask = nullptr; - ggml_tensor * conv_pad_mask = nullptr; - bool use_flash = true; + ggml_tensor * attn_pad_mask = nullptr; + ggml_tensor * conv_pad_mask = nullptr; + bool use_flash = true; + bool direct_depthwise = false; + bool bounded_depthwise = false; }; // LayerNorm with kLayerNormEps. `beta` is optional (may be nullptr). @@ -85,7 +87,9 @@ ggml_tensor * fsmn_branch(ggml_context * ctx, ggml_tensor * v_pre, ggml_tensor * fsmn_w, int kernel, - ggml_tensor * conv_pad_mask = nullptr); + ggml_tensor * conv_pad_mask, + bool direct_depthwise, + bool bounded_depthwise); // SAN-M attention sub-block: fused QKV, FSMN parallel branch on V, // SDPA over the QKV split. Returns ne=[d_model, T, B] (post-projection diff --git a/src/transcribe-batch-util.cpp b/src/transcribe-batch-util.cpp index 11372a9a..7b73db8d 100644 --- a/src/transcribe-batch-util.cpp +++ b/src/transcribe-batch-util.cpp @@ -197,6 +197,31 @@ transcribe_status decode_batch_slices(transcribe_session * session, return TRANSCRIBE_OK; } +transcribe_status decode_batch_id_slices( + transcribe_session * session, + int n, + const int32_t * host_buf, + std::size_t utt_elems, + int64_t total_encode_us, + int64_t total_mel_us, + const std::function & decode_fn) { + const int64_t enc_per_utt = total_encode_us / std::max(1, n); + const int64_t mel_per_utt = total_mel_us / std::max(1, n); + for (int b = 0; b < n; ++b) { + if (session->poll_abort()) { + return TRANSCRIBE_ERR_ABORTED; + } + session->clear_result(); + const int32_t * slice = host_buf + static_cast(b) * utt_elems; + const transcribe_status st = decode_fn(b, slice); + auto rs = session->capture_result(st); + rs.t_mel_us = mel_per_utt; + rs.t_encode_us = enc_per_utt; + session->batch_results.push_back(std::move(rs)); + } + return TRANSCRIBE_OK; +} + transcribe_status run_batched_encdec_step_loop(transcribe_session * session, ggml_backend_sched_t sched, const EncDecRebuildFn & rebuild, diff --git a/src/transcribe-batch-util.h b/src/transcribe-batch-util.h index 9be59eed..30abb506 100644 --- a/src/transcribe-batch-util.h +++ b/src/transcribe-batch-util.h @@ -102,6 +102,16 @@ transcribe_status decode_batch_slices(transcribe_session * session, int64_t total_mel_us, const std::function & decode_fn); +// Integer counterpart for device-side CTC argmax outputs. +transcribe_status decode_batch_id_slices( + transcribe_session * session, + int n, + const int32_t * host_buf, + std::size_t utt_elems, + int64_t total_encode_us, + int64_t total_mel_us, + const std::function & decode_fn); + // --------------------------------------------------------------------------- // Batched encoder-decoder greedy step loop (cohere / canary / moonshine) // From 9d3e8456601d57bd921c8facb9b406bb996849cc Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 07:29:40 +0800 Subject: [PATCH 12/17] moonshine mem reduction --- src/arch/moonshine/model.cpp | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/src/arch/moonshine/model.cpp b/src/arch/moonshine/model.cpp index 9f3de4a4..4b687652 100644 --- a/src/arch/moonshine/model.cpp +++ b/src/arch/moonshine/model.cpp @@ -194,6 +194,19 @@ namespace { constexpr const char k_default_variant[] = "moonshine"; +// Upstream recommends roughly 6.5 generated tokens per second. Keep generous +// short-clip headroom while avoiding a full 194-position self-KV allocation +// for every utterance. +int decode_generation_budget(int n_samples, int model_max) { + constexpr int64_t k_native_sr_hz = 16000; + constexpr int64_t k_budget_num = 13; + constexpr int64_t k_budget_den = 2; + constexpr int64_t k_budget_floor = 24; + const int duration_budget = static_cast( + static_cast(n_samples) * k_budget_num / (k_budget_den * k_native_sr_hz) + k_budget_floor); + return model_max > 0 ? std::min(duration_budget, model_max) : duration_budget; +} + extern transcribe_status load(Loader &, const transcribe_model_load_params *, transcribe_model **); extern transcribe_status init_context(transcribe_model *, const transcribe_session_params *, transcribe_session **); extern transcribe_status run(transcribe_session *, const float *, int, const transcribe_run_params *); @@ -437,18 +450,19 @@ transcribe_status run(transcribe_session * session, ggml_backend_tensor_get(eb.out, cc->enc_host.data(), 0, cc->enc_host.size() * sizeof(float)); // ----- KV cache init ----- + const int decode_cap = decode_generation_budget(n_samples, hp.dec_max_position_embeddings); { - if (cc->kv_cache.buffer != nullptr && cc->kv_cache.T_enc != T_enc) { + if (cc->kv_cache.buffer != nullptr && (cc->kv_cache.T_enc != T_enc || cc->kv_cache.n_ctx != decode_cap)) { cc->kv_cache.free(); } if (cc->kv_cache.buffer == nullptr) { - const int n_ctx = hp.dec_max_position_embeddings > 0 ? hp.dec_max_position_embeddings : 512; ggml_type cache_type = resolved_kv; // Default the cache to F32 to match moonshine's reference regime. if (cache_type == GGML_TYPE_COUNT) { cache_type = GGML_TYPE_F32; } - if (!kv_cache_init(cc->kv_cache, cm->plan.primary, n_ctx, T_enc, d_model, hp.dec_n_layers, cache_type)) { + if (!kv_cache_init(cc->kv_cache, cm->plan.primary, decode_cap, T_enc, d_model, hp.dec_n_layers, + cache_type)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moonshine run: KV cache allocation failed — out of memory."); return TRANSCRIBE_ERR_OOM; @@ -492,7 +506,7 @@ transcribe_status run(transcribe_session * session, const int64_t t_decode_start = ggml_time_us(); const int decoder_start = hp.decoder_start_token_id; // 1 const int eos = hp.eos_token_id; // 2 - const int max_pos = hp.dec_max_position_embeddings; + const int max_pos = decode_cap; std::vector generated_ids; int next_token = -1; @@ -903,7 +917,13 @@ transcribe_status run_batch(transcribe_session * session, const int n_layer = hp.dec_n_layers; const int decoder_start = hp.decoder_start_token_id; const int32_t eos = hp.eos_token_id; - const int max_pos = hp.dec_max_position_embeddings; + int max_samples = 0; + for (int b = 0; b < n; ++b) { + if (pcm[b] != nullptr && n_samples[b] > 0) { + max_samples = std::max(max_samples, n_samples[b]); + } + } + const int max_pos = decode_generation_budget(max_samples, hp.dec_max_position_embeddings); // ----- Serial per-utterance encoder (no mel; raw PCM) ----- std::vector valid(n, 0); From cd3ebb89c3d5263ba75be2e1c7b2f81158a1c469 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 07:29:45 +0800 Subject: [PATCH 13/17] granite5 mem reduction --- src/arch/granite5_ctc/decoder.cpp | 22 +++++ src/arch/granite5_ctc/decoder.h | 7 ++ src/arch/granite5_ctc/encoder.cpp | 136 ++++++++++++++++++++++++--- src/arch/granite5_ctc/encoder.h | 7 +- src/arch/granite5_ctc/granite5_ctc.h | 6 +- src/arch/granite5_ctc/model.cpp | 89 +++++++++++++----- 6 files changed, 229 insertions(+), 38 deletions(-) diff --git a/src/arch/granite5_ctc/decoder.cpp b/src/arch/granite5_ctc/decoder.cpp index 51cf8fd3..52c58ad8 100644 --- a/src/arch/granite5_ctc/decoder.cpp +++ b/src/arch/granite5_ctc/decoder.cpp @@ -66,6 +66,28 @@ void collapse_whitespace(std::string & s) { } // namespace +void ctc_greedy_collapse_ids(const int32_t * ids, + const float * probs, + int t_len, + int blank_id, + std::vector & out_tokens) { + out_tokens.clear(); + if (ids == nullptr || probs == nullptr || t_len <= 0) { + return; + } + int prev = -1; + for (int t = 0; t < t_len; ++t) { + const int label = ids[t]; + if (label == prev) { + continue; + } + prev = label; + if (label != blank_id) { + out_tokens.push_back({ label, probs[t], t }); + } + } +} + void ctc_greedy_collapse(const float * logits, int t_len, int vocab, int blank_id, std::vector & out_tokens) { out_tokens.clear(); if (logits == nullptr || t_len <= 0 || vocab <= 0) { diff --git a/src/arch/granite5_ctc/decoder.h b/src/arch/granite5_ctc/decoder.h index 1d078227..b15290db 100644 --- a/src/arch/granite5_ctc/decoder.h +++ b/src/arch/granite5_ctc/decoder.h @@ -44,6 +44,13 @@ struct CtcToken { // [T, vocab] host order (i.e. ggml ne = [vocab, T] read contiguously). void ctc_greedy_collapse(const float * logits, int t_len, int vocab, int blank_id, std::vector & out_tokens); +// Collapse device-computed frame argmax IDs and winning-class probabilities. +void ctc_greedy_collapse_ids(const int32_t * ids, + const float * probs, + int t_len, + int blank_id, + std::vector & out_tokens); + // Milliseconds per encoder frame. Every subsampling block halves the // rate on top of the frontend's hop and frame stacking: // hop * stack * 2^len(subsample_layers) / sample_rate diff --git a/src/arch/granite5_ctc/encoder.cpp b/src/arch/granite5_ctc/encoder.cpp index 79d1efe6..f7e611f8 100644 --- a/src/arch/granite5_ctc/encoder.cpp +++ b/src/arch/granite5_ctc/encoder.cpp @@ -48,6 +48,98 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, ggml_tensor * w, ggml_ return y; } +// Concatenate a potentially long chunk list as a balanced tree so graph +// traversal depth stays logarithmic. +ggml_tensor * concat_columns(ggml_context * ctx, std::vector chunks) { + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/1) : chunks[i]); + } + chunks.swap(next); + } + return chunks.front(); +} + +// The self-conditioning CTC vocabulary projection is frame-local. Chunk it +// to avoid retaining full [vocab, time] logits and softmax tensors. +ggml_tensor * self_condition_chunked(ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * proj_w, + ggml_tensor * proj_b, + ggml_tensor * bypass_w, + ggml_tensor * bypass_b) { + constexpr int64_t kChunk = 128; + const int64_t hidden = x->ne[0]; + const int64_t cols = x->ne[1] * x->ne[2]; + ggml_tensor * flat = ggml_reshape_2d(ctx, x, hidden, cols); + std::vector chunks; + chunks.reserve(static_cast((cols + kChunk - 1) / kChunk)); + for (int64_t col0 = 0; col0 < cols; col0 += kChunk) { + const int64_t n = std::min(kChunk, cols - col0); + ggml_tensor * part = ggml_view_2d(ctx, flat, hidden, n, flat->nb[1], col0 * flat->nb[1]); + ggml_tensor * logits = linear(ctx, part, proj_w, proj_b); + chunks.push_back(linear(ctx, ggml_soft_max(ctx, logits), bypass_w, bypass_b)); + } + ggml_tensor * joined = concat_columns(ctx, std::move(chunks)); + return ggml_reshape_3d(ctx, joined, hidden, x->ne[1], x->ne[2]); +} + +struct ChunkedCtc { + ggml_tensor * ids = nullptr; + ggml_tensor * probs = nullptr; +}; + +// Produce greedy labels and their softmax probabilities in bounded chunks. +// get_rows selects each winning class; multiplying by a chunk-local identity +// extracts the diagonal without a backend-specific gather-elements op. +ChunkedCtc ctc_argmax_chunked(ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * proj_w, + ggml_tensor * proj_b, + std::vector & diag_masks) { + constexpr int64_t kChunk = 256; + const int64_t hidden = x->ne[0]; + const int64_t cols = x->ne[1] * x->ne[2]; + ggml_tensor * flat = ggml_reshape_2d(ctx, x, hidden, cols); + std::vector id_chunks; + std::vector prob_chunks; + id_chunks.reserve(static_cast((cols + kChunk - 1) / kChunk)); + prob_chunks.reserve(id_chunks.capacity()); + diag_masks.reserve(id_chunks.capacity()); + for (int64_t col0 = 0; col0 < cols; col0 += kChunk) { + const int64_t n = std::min(kChunk, cols - col0); + ggml_tensor * part = ggml_view_2d(ctx, flat, hidden, n, flat->nb[1], col0 * flat->nb[1]); + ggml_tensor * logits = linear(ctx, part, proj_w, proj_b); + ggml_tensor * ids = ggml_argmax(ctx, logits); + id_chunks.push_back(ids); + + ggml_tensor * probs_t = ggml_cont(ctx, ggml_transpose(ctx, ggml_soft_max(ctx, logits))); + ggml_tensor * picked = ggml_get_rows(ctx, probs_t, ids); // [frame, selected frame] + ggml_tensor * diag = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n, n); + ggml_set_input(diag); + diag_masks.push_back(diag); + prob_chunks.push_back(ggml_reshape_1d(ctx, ggml_sum_rows(ctx, ggml_mul(ctx, picked, diag)), n)); + } + auto concat_vectors = [ctx](std::vector chunks) { + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/0) : + chunks[i]); + } + chunks.swap(next); + } + return chunks.front(); + }; + return { + ggml_reshape_2d(ctx, concat_vectors(std::move(id_chunks)), x->ne[1], x->ne[2]), + ggml_reshape_2d(ctx, concat_vectors(std::move(prob_chunks)), x->ne[1], x->ne[2]), + }; +} + } // namespace // Host-side frontend. @@ -583,13 +675,18 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // of block (self_cond_layer - 1), using the SAME projection as // the final CTC head (tied weights). if (i + 1 == hp.enc_self_cond_layer) { - ggml_tensor * mid_logits = linear(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b); - record("enc.ctc.mid_logits", mid_logits); - - ggml_tensor * mid_soft = ggml_soft_max(ctx, mid_logits); - ggml_tensor * injection = linear(ctx, mid_soft, weights.enc_top.ctc_bypass_w, weights.enc_top.ctc_bypass_b); - record("enc.ctc.mid_injection", injection); - + ggml_tensor * injection; + if (transcribe::debug::enabled()) { + ggml_tensor * mid_logits = linear(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b); + record("enc.ctc.mid_logits", mid_logits); + + ggml_tensor * mid_soft = ggml_soft_max(ctx, mid_logits); + injection = linear(ctx, mid_soft, weights.enc_top.ctc_bypass_w, weights.enc_top.ctc_bypass_b); + record("enc.ctc.mid_injection", injection); + } else { + injection = self_condition_chunked(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b, + weights.enc_top.ctc_bypass_w, weights.enc_top.ctc_bypass_b); + } x = ggml_add(ctx, x, injection); } } @@ -603,17 +700,32 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_set_output(eb.out); eb.dump_list.emplace_back("enc.out", eb.out); - // Tied CTC head: the same projection as the mid-layer one. - eb.ctc_logits = linear(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b); - record("enc.ctc_logits", eb.ctc_logits); - ggml_set_output(eb.ctc_logits); + // Tied CTC head: retain complete logits only for numerical validation. + // Normal greedy inference downloads one int32 label per frame instead. + if (transcribe::debug::enabled()) { + eb.ctc_logits = linear(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b); + record("enc.ctc_logits", eb.ctc_logits); + ggml_set_output(eb.ctc_logits); + } else { + const ChunkedCtc ctc = + ctc_argmax_chunked(ctx, x, weights.enc_top.ctc_proj_w, weights.enc_top.ctc_proj_b, eb.ctc_diag_masks); + eb.ctc_ids = ctc.ids; + eb.ctc_probs = ctc.probs; + named(eb.ctc_ids, "enc.ctc_ids"); + named(eb.ctc_probs, "enc.ctc_probs"); + ggml_set_output(eb.ctc_ids); + ggml_set_output(eb.ctc_probs); + } eb.graph = ggml_new_graph_custom(ctx, /*size=*/16384, /*grads=*/false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "granite5_ctc encoder: ggml_new_graph_custom failed"); return eb; } - ggml_build_forward_expand(eb.graph, eb.ctc_logits); + ggml_build_forward_expand(eb.graph, eb.ctc_logits != nullptr ? eb.ctc_logits : eb.ctc_ids); + if (eb.ctc_probs != nullptr) { + ggml_build_forward_expand(eb.graph, eb.ctc_probs); + } ggml_build_forward_expand(eb.graph, eb.out); return eb; diff --git a/src/arch/granite5_ctc/encoder.h b/src/arch/granite5_ctc/encoder.h index 6da1b816..184f3066 100644 --- a/src/arch/granite5_ctc/encoder.h +++ b/src/arch/granite5_ctc/encoder.h @@ -144,8 +144,11 @@ struct EncoderBuild { std::vector block_stage; // block index -> stages[] index // Graph outputs. - ggml_tensor * out = nullptr; // [hidden, T_out] ("enc.out") - ggml_tensor * ctc_logits = nullptr; // [vocab, T_out] ("enc.ctc_logits") + ggml_tensor * out = nullptr; // [hidden, T_out, B] ("enc.out") + ggml_tensor * ctc_logits = nullptr; // [vocab, T_out, B], debug validation only + ggml_tensor * ctc_ids = nullptr; // [T_out, B] i32, normal greedy path + ggml_tensor * ctc_probs = nullptr; // [T_out, B] f32 winning-class probabilities + std::vector ctc_diag_masks; // chunk-local identity inputs ggml_cgraph * graph = nullptr; diff --git a/src/arch/granite5_ctc/granite5_ctc.h b/src/arch/granite5_ctc/granite5_ctc.h index 072faa41..d9663055 100644 --- a/src/arch/granite5_ctc/granite5_ctc.h +++ b/src/arch/granite5_ctc/granite5_ctc.h @@ -63,8 +63,10 @@ struct Granite5CtcModel final : public transcribe_model { struct Granite5CtcSession final : public transcribe_session { // Host scratch, reused across runs. - std::vector feats_buf; // [T_enc, input_dim] stacked frontend output - std::vector logits_buf; // [T_out, vocab] CTC logits + std::vector feats_buf; // [T_enc, input_dim] stacked frontend output + std::vector logits_buf; // [T_out, vocab] CTC logits (debug only) + std::vector ctc_ids; // [T_out] greedy labels + std::vector ctc_probs; // [T_out] winning-class probabilities Granite5CtcSession() = default; ~Granite5CtcSession() override; diff --git a/src/arch/granite5_ctc/model.cpp b/src/arch/granite5_ctc/model.cpp index 2d80d597..0b106c12 100644 --- a/src/arch/granite5_ctc/model.cpp +++ b/src/arch/granite5_ctc/model.cpp @@ -293,7 +293,8 @@ transcribe_status run_encoder(Granite5CtcSession * cc, } out_eb = build_encoder_graph(cc->compute_ctx, cm->weights, hp, T_max, real_lens); - if (out_eb.graph == nullptr || out_eb.out == nullptr || out_eb.ctc_logits == nullptr) { + if (out_eb.graph == nullptr || out_eb.out == nullptr || + (out_eb.ctc_logits == nullptr && out_eb.ctc_ids == nullptr)) { return TRANSCRIBE_ERR_GGUF; } @@ -348,6 +349,15 @@ transcribe_status run_encoder(Granite5CtcSession * cc, } } + for (ggml_tensor * diag : out_eb.ctc_diag_masks) { + const int64_t n = diag->ne[0]; + std::vector identity(static_cast(n * n), 0.0f); + for (int64_t i = 0; i < n; ++i) { + identity[static_cast(i * n + i)] = 1.0f; + } + ggml_backend_tensor_set(diag, identity.data(), 0, identity.size() * sizeof(float)); + } + transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); @@ -396,6 +406,19 @@ void decode_and_populate(Granite5CtcSession * cc, cc->t_decode_us = ggml_time_us() - t_dec_start; } +void decode_ids_and_populate(Granite5CtcSession * cc, + Granite5CtcModel * cm, + const int32_t * ids, + const float * probs, + int t_valid, + int64_t clip_ms) { + const int64_t t_dec_start = ggml_time_us(); + std::vector toks; + ctc_greedy_collapse_ids(ids, probs, t_valid, cm->hparams.blank_id, toks); + build_result(*cc, cm->tok, cm->hparams, toks, clip_ms); + cc->t_decode_us = ggml_time_us() - t_dec_start; +} + int64_t clip_ms_for(const Granite5CtcHParams & hp, int n_samples) { return (hp.fe_sample_rate > 0) ? (static_cast(n_samples) * 1000 / hp.fe_sample_rate) : 0; } @@ -420,13 +443,20 @@ transcribe_status encode_and_decode(Granite5CtcSession * cc, Granite5CtcModel * return st; } - const int t_out = static_cast(eb.ctc_logits->ne[1]); - const int vocab = static_cast(eb.ctc_logits->ne[0]); - cc->logits_buf.assign(static_cast(t_out) * vocab, 0.0f); - ggml_backend_tensor_get(eb.ctc_logits, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); - - decode_and_populate(cc, cm, cc->logits_buf.data(), t_out, vocab, clip_ms_for(hp, n_samples), - /*utt_index=*/-1); + const int t_out = eb.t_out; + if (eb.ctc_logits != nullptr) { + const int vocab = static_cast(eb.ctc_logits->ne[0]); + cc->logits_buf.assign(static_cast(t_out) * vocab, 0.0f); + ggml_backend_tensor_get(eb.ctc_logits, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); + decode_and_populate(cc, cm, cc->logits_buf.data(), t_out, vocab, clip_ms_for(hp, n_samples), + /*utt_index=*/-1); + } else { + cc->ctc_ids.assign(static_cast(t_out), 0); + cc->ctc_probs.assign(static_cast(t_out), 0.0f); + ggml_backend_tensor_get(eb.ctc_ids, cc->ctc_ids.data(), 0, cc->ctc_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_get(eb.ctc_probs, cc->ctc_probs.data(), 0, cc->ctc_probs.size() * sizeof(float)); + decode_ids_and_populate(cc, cm, cc->ctc_ids.data(), cc->ctc_probs.data(), t_out, clip_ms_for(hp, n_samples)); + } return TRANSCRIBE_OK; } @@ -525,20 +555,35 @@ transcribe_status run_batch(transcribe_session * session, EncoderBuild eb{}; if (auto st = run_encoder(cc, cm, T_max, lens, eb); st == TRANSCRIBE_OK) { - const int t_out = static_cast(eb.ctc_logits->ne[1]); - const int vocab = static_cast(eb.ctc_logits->ne[0]); - const size_t utt_elems = static_cast(t_out) * vocab; - // Full read then host-slice: non-zero-offset backend reads are - // not reliable across every backend. - cc->logits_buf.assign(utt_elems * static_cast(n), 0.0f); - ggml_backend_tensor_get(eb.ctc_logits, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); - - return transcribe::decode_batch_slices( - cc, n, cc->logits_buf.data(), utt_elems, cc->t_encode_us, total_mel_us, - [&](int b, const float * logits_b) -> transcribe_status { - const int t_valid = std::min(eb.real_lens_out[static_cast(b)], t_out); - decode_and_populate(cc, cm, logits_b, t_valid, vocab, clip_ms_for(cm->hparams, n_samples[b]), - /*utt_index=*/b); + const int t_out = eb.t_out; + if (eb.ctc_logits != nullptr) { + const int vocab = static_cast(eb.ctc_logits->ne[0]); + const size_t utt_elems = static_cast(t_out) * vocab; + // Full read then host-slice: non-zero-offset backend reads are + // not reliable across every backend. + cc->logits_buf.assign(utt_elems * static_cast(n), 0.0f); + ggml_backend_tensor_get(eb.ctc_logits, cc->logits_buf.data(), 0, cc->logits_buf.size() * sizeof(float)); + + return transcribe::decode_batch_slices( + cc, n, cc->logits_buf.data(), utt_elems, cc->t_encode_us, total_mel_us, + [&](int b, const float * logits_b) -> transcribe_status { + const int t_valid = std::min(eb.real_lens_out[static_cast(b)], t_out); + decode_and_populate(cc, cm, logits_b, t_valid, vocab, clip_ms_for(cm->hparams, n_samples[b]), + /*utt_index=*/b); + return TRANSCRIBE_OK; + }); + } + + cc->ctc_ids.assign(static_cast(t_out) * static_cast(n), 0); + cc->ctc_probs.assign(static_cast(t_out) * static_cast(n), 0.0f); + ggml_backend_tensor_get(eb.ctc_ids, cc->ctc_ids.data(), 0, cc->ctc_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_get(eb.ctc_probs, cc->ctc_probs.data(), 0, cc->ctc_probs.size() * sizeof(float)); + return transcribe::decode_batch_id_slices( + cc, n, cc->ctc_ids.data(), static_cast(t_out), cc->t_encode_us, total_mel_us, + [&](int b, const int32_t * ids_b) -> transcribe_status { + const int t_valid = std::min(eb.real_lens_out[static_cast(b)], t_out); + const float * probs_b = cc->ctc_probs.data() + static_cast(b) * t_out; + decode_ids_and_populate(cc, cm, ids_b, probs_b, t_valid, clip_ms_for(cm->hparams, n_samples[b])); return TRANSCRIBE_OK; }); } From 4e0f8b302fc66be92d75eb3bbe44152bb3556a1c Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 17:05:04 +0800 Subject: [PATCH 14/17] canary mem reduction --- src/arch/canary/canary.h | 4 ---- src/arch/canary/encoder.cpp | 10 ++++++---- src/arch/canary/model.cpp | 29 ----------------------------- src/conformer/conformer.cpp | 14 ++++++++++---- src/conformer/conformer.h | 1 + 5 files changed, 17 insertions(+), 41 deletions(-) diff --git a/src/arch/canary/canary.h b/src/arch/canary/canary.h index 27bac011..7e898723 100644 --- a/src/arch/canary/canary.h +++ b/src/arch/canary/canary.h @@ -110,10 +110,6 @@ struct CanaryModel final : public transcribe_model { ggml_context * bn_fused_ctx = nullptr; ggml_backend_buffer_t bn_fused_buffer = nullptr; - // CPU-only F16 -> F32 promotion buffer for conformer 1x1 pointwise convs. - ggml_context * conv_pw_f32_ctx = nullptr; - ggml_backend_buffer_t conv_pw_f32_buffer = nullptr; - std::optional mel; CanaryModel() = default; diff --git a/src/arch/canary/encoder.cpp b/src/arch/canary/encoder.cpp index 744c7194..fc3611a1 100644 --- a/src/arch/canary/encoder.cpp +++ b/src/arch/canary/encoder.cpp @@ -102,10 +102,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, bool use_flash, const char * backend_name) { conf::ConvPolicy policy{}; - policy.direct_pw = conf::detect_direct_pw(backend_name); - const bool direct_dw = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_block = direct_dw; - policy.direct_dw_in_pre_encode = false; // parakeet-style: im2col here + policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.promote_pw_in_graph = backend_name != nullptr && std::strstr(backend_name, "CPU") != nullptr; + const bool direct_dw = detect_direct_dw_in_block(backend_name); + policy.direct_dw_in_block = direct_dw; + policy.direct_dw_in_pre_encode = false; // parakeet-style: im2col here + policy.pre_encode_dw_time_chunk = 256; EncoderBuild eb{}; diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index 6b556357..6295a1a5 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -182,14 +182,6 @@ CanaryModel::~CanaryModel() { safe_buffer_free(bn_fused_buffer); bn_fused_buffer = nullptr; } - if (conv_pw_f32_ctx != nullptr) { - ggml_free(conv_pw_f32_ctx); - conv_pw_f32_ctx = nullptr; - } - if (conv_pw_f32_buffer != nullptr) { - safe_buffer_free(conv_pw_f32_buffer); - conv_pw_f32_buffer = nullptr; - } if (ctx_meta != nullptr) { ggml_free(ctx_meta); ctx_meta = nullptr; @@ -322,23 +314,6 @@ transcribe_status fuse_batch_norm(CanaryModel & m) { return TRANSCRIBE_OK; } -// On CPU primary backend, dequantize 1x1 conformer pointwise convs -// from F16 to F32. Same rationale as parakeet/cohere. -transcribe_status promote_conv_pw_to_f32_on_cpu(CanaryModel & m) { - std::vector slots; - slots.reserve(m.weights.blocks.size() * 2); - for (auto & b : m.weights.blocks) { - if (b.conv_pw1_w != nullptr && b.conv_pw1_w->type == GGML_TYPE_F16) { - slots.push_back({ &b.conv_pw1_w, b.conv_pw1_w }); - } - if (b.conv_pw2_w != nullptr && b.conv_pw2_w->type == GGML_TYPE_F16) { - slots.push_back({ &b.conv_pw2_w, b.conv_pw2_w }); - } - } - return load_common::promote_conv_pw_f16_to_f32_on_cpu(m.plan, slots, "canary", &m.conv_pw_f32_ctx, - &m.conv_pw_f32_buffer); -} - constexpr const char k_default_variant[] = "canary"; extern transcribe_status load(Loader &, const transcribe_model_load_params *, transcribe_model **); @@ -505,10 +480,6 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par if (const transcribe_status st = fuse_batch_norm(*m); st != TRANSCRIBE_OK) { return st; } - if (const transcribe_status st = promote_conv_pw_to_f32_on_cpu(*m); st != TRANSCRIBE_OK) { - return st; - } - m->t_load_us = ggml_time_us() - t_load_start; *out_model = m.release(); return TRANSCRIBE_OK; diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 704b9b0a..74fcfe1e 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -366,8 +366,11 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & // causal_lm.cpp for the CUDA COMPUTE_16F saturation rationale). { ggml_tensor * pw1 = ggml_reshape_2d(ctx, b.conv_pw1_w, d_model, 2 * d_model); - x = ggml_mul_mat(ctx, pw1, x); // [2*d_model, T, B] - if (b.conv_pw1_w->type == GGML_TYPE_F16) { + if (policy.promote_pw_in_graph && pw1->type == GGML_TYPE_F16) { + pw1 = ggml_cast(ctx, pw1, GGML_TYPE_F32); + } + x = ggml_mul_mat(ctx, pw1, x); // [2*d_model, T, B] + if (pw1->type == GGML_TYPE_F16) { ggml_mul_mat_set_prec(x, GGML_PREC_F32); } if (b.conv_pw1_b != nullptr) { @@ -518,8 +521,11 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & // Pointwise conv 2 as direct mul_mat in [d_model, T] layout. See // the pw1 comment above for the F16 / CUDA COMPUTE_16F rationale. ggml_tensor * pw2 = ggml_reshape_2d(ctx, b.conv_pw2_w, d_model, d_model); - x = ggml_mul_mat(ctx, pw2, x); - if (b.conv_pw2_w->type == GGML_TYPE_F16) { + if (policy.promote_pw_in_graph && pw2->type == GGML_TYPE_F16) { + pw2 = ggml_cast(ctx, pw2, GGML_TYPE_F32); + } + x = ggml_mul_mat(ctx, pw2, x); + if (pw2->type == GGML_TYPE_F16) { ggml_mul_mat_set_prec(x, GGML_PREC_F32); } if (b.conv_pw2_b != nullptr) { diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index 3df9afb4..4c7163f8 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -108,6 +108,7 @@ struct BlockView { // paths except for direct_pw; families opt into direct pre-encode ops. struct ConvPolicy { bool direct_pw = true; + bool promote_pw_in_graph = false; bool direct_conv0_in_pre_encode = false; bool direct_dw_in_block = false; bool direct_dw_in_pre_encode = false; From f8f9ffa646101eacd0781a78d14072fd2d43a840 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 17:05:04 +0800 Subject: [PATCH 15/17] cohere mem reduction --- src/arch/cohere/cohere.h | 6 ------ src/arch/cohere/encoder.cpp | 1 + src/arch/cohere/model.cpp | 32 -------------------------------- 3 files changed, 1 insertion(+), 38 deletions(-) diff --git a/src/arch/cohere/cohere.h b/src/arch/cohere/cohere.h index 0a7e9860..5794bda7 100644 --- a/src/arch/cohere/cohere.h +++ b/src/arch/cohere/cohere.h @@ -138,12 +138,6 @@ struct CohereModel final : public transcribe_model { ggml_context * bn_fused_ctx = nullptr; ggml_backend_buffer_t bn_fused_buffer = nullptr; - // On CPU primary backend, the conformer 1×1 pointwise conv weights - // are dequantized F16->F32 at load time (Zen 2 has no native F16 - // compute). Tensors live here; CohereBlock slots point at them. - ggml_context * conv_pw_f32_ctx = nullptr; - ggml_backend_buffer_t conv_pw_f32_buffer = nullptr; - std::optional mel; CohereModel() = default; diff --git a/src/arch/cohere/encoder.cpp b/src/arch/cohere/encoder.cpp index 75ebc721..a0b53d0b 100644 --- a/src/arch/cohere/encoder.cpp +++ b/src/arch/cohere/encoder.cpp @@ -113,6 +113,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const char * backend_name) { conf::ConvPolicy policy{}; policy.direct_pw = conf::detect_direct_pw(backend_name); + policy.promote_pw_in_graph = backend_name != nullptr && std::strstr(backend_name, "CPU") != nullptr; const bool direct_dw = detect_direct_dw(backend_name); policy.direct_dw_in_block = direct_dw; policy.direct_dw_in_pre_encode = direct_dw; diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index a305756f..846b8d40 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -192,14 +192,6 @@ CohereModel::~CohereModel() { safe_buffer_free(bn_fused_buffer); bn_fused_buffer = nullptr; } - if (conv_pw_f32_ctx != nullptr) { - ggml_free(conv_pw_f32_ctx); - conv_pw_f32_ctx = nullptr; - } - if (conv_pw_f32_buffer != nullptr) { - safe_buffer_free(conv_pw_f32_buffer); - conv_pw_f32_buffer = nullptr; - } if (ctx_meta != nullptr) { ggml_free(ctx_meta); ctx_meta = nullptr; @@ -395,25 +387,6 @@ transcribe_status fuse_encoder_q_bias(CohereModel & m) { return TRANSCRIBE_OK; } -// On a CPU primary backend, dequantize the conformer 1×1 pointwise conv -// weights (pw1, pw2) from F16 back to F32: Zen 2 (and anything else without -// native F16 compute) pays an F16->F32 upconvert per matmul that outweighs -// the bandwidth win. GPU backends skip this and keep the F16 weights. -transcribe_status promote_conv_pw_to_f32_on_cpu(CohereModel & m) { - std::vector slots; - slots.reserve(m.weights.blocks.size() * 2); - for (auto & b : m.weights.blocks) { - if (b.conv_pw1_w != nullptr && b.conv_pw1_w->type == GGML_TYPE_F16) { - slots.push_back({ &b.conv_pw1_w, b.conv_pw1_w }); - } - if (b.conv_pw2_w != nullptr && b.conv_pw2_w->type == GGML_TYPE_F16) { - slots.push_back({ &b.conv_pw2_w, b.conv_pw2_w }); - } - } - return load_common::promote_conv_pw_f16_to_f32_on_cpu(m.plan, slots, "cohere", &m.conv_pw_f32_ctx, - &m.conv_pw_f32_buffer); -} - constexpr const char k_default_variant[] = "cohere-asr"; // Forward declarations for the Arch trait below. @@ -595,11 +568,6 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par return st; } - // CPU only: dequantize conv pointwise weights to F32 (see function doc). - if (const transcribe_status st = promote_conv_pw_to_f32_on_cpu(*m); st != TRANSCRIBE_OK) { - return st; - } - m->t_load_us = ggml_time_us() - t_load_start; *out_model = m.release(); return TRANSCRIBE_OK; From 56d6c9114678cd2e878e5167c5a77c989fab78e1 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 17:05:04 +0800 Subject: [PATCH 16/17] granite nar mem reduction --- src/arch/granite_nar/encoder.cpp | 68 ++++++++++++++++++++++++++------ src/arch/granite_nar/encoder.h | 3 +- src/arch/granite_nar/model.cpp | 3 +- 3 files changed, 61 insertions(+), 13 deletions(-) diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index 5470ff71..2eaad854 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/src/arch/granite_nar/encoder.cpp @@ -172,7 +172,8 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * bn_fused_scale, ggml_tensor * bn_fused_bias, int conv_kernel, - int inner_dim) { + int inner_dim, + bool direct_depthwise) { const int64_t d_model = x->ne[0]; const int64_t T = x->ne[1]; @@ -187,13 +188,51 @@ ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * value = ggml_view_2d(ctx, x, inner_dim, T, x->nb[1], inner_dim * ggml_element_size(x)); x = ggml_mul(ctx, gate, ggml_sigmoid(ctx, value)); } - x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + const int padding = (conv_kernel - 1) / 2; - x = transcribe::conformer::conv_1d_dw_f32(ctx, b.conv_depthwise_w, x, - /*stride=*/1, /*padding=*/padding, /*dilation=*/1); - x = transcribe::conformer::fused_batch_norm(ctx, x, bn_fused_scale, bn_fused_bias); - x = ggml_silu(ctx, x); - x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + if (direct_depthwise) { + ggml_tensor * kernel = ggml_reshape_4d(ctx, b.conv_depthwise_w, conv_kernel, 1, 1, inner_dim); + ggml_tensor * data = ggml_reshape_4d(ctx, x, T, 1, inner_dim, 1); + x = transcribe::conformer::conv_2d_dw_direct_f32(ctx, kernel, data, + /*s0=*/1, /*s1=*/1, + /*p0=*/padding, /*p1=*/0, + /*d0=*/1, /*d1=*/1); + x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[2], x->ne[3]); + } else { + constexpr int64_t kTimeChunk = 256; + std::vector chunks; + chunks.reserve(static_cast((T + kTimeChunk - 1) / kTimeChunk)); + for (int64_t out0 = 0; out0 < T; out0 += kTimeChunk) { + const int64_t n_out = std::min(kTimeChunk, T - out0); + const int64_t src0 = out0 - padding; + const int64_t src1 = src0 + n_out + conv_kernel - 1; + const int64_t view0 = std::max(0, src0); + const int64_t view1 = std::min(T, src1); + const int64_t pad_left = view0 - src0; + const int64_t pad_right = src1 - view1; + ggml_tensor * input = + ggml_view_3d(ctx, x, view1 - view0, inner_dim, 1, x->nb[1], x->nb[2], view0 * x->nb[0]); + input = ggml_pad_ext(ctx, input, pad_left, pad_right, 0, 0, 0, 0, 0, 0); + chunks.push_back(transcribe::conformer::conv_1d_dw_f32(ctx, b.conv_depthwise_w, input, + /*stride=*/1, /*padding=*/0, + /*dilation=*/1)); + } + while (chunks.size() > 1) { + std::vector next; + next.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + next.push_back(i + 1 < chunks.size() ? ggml_concat(ctx, chunks[i], chunks[i + 1], /*dim=*/0) : + chunks[i]); + } + chunks.swap(next); + } + x = chunks.front(); + } + + x = transcribe::conformer::fused_batch_norm(ctx, x, bn_fused_scale, bn_fused_bias); + x = ggml_silu(ctx, x); + x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); { ggml_tensor * pw2 = ggml_reshape_2d(ctx, b.conv_pointwise2_w, inner_dim, d_model); x = ggml_mul_mat(ctx, pw2, x); @@ -208,8 +247,15 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const GraniteNarWeights & weights, const GraniteNarHParams & hp, int T_enc, - bool /*use_flash*/) { + bool /*use_flash*/, + const char * backend_name) { EncoderBuild eb{}; + + const bool backend_direct = backend_name != nullptr && (std::strstr(backend_name, "Vulkan") != nullptr || + std::strstr(backend_name, "CUDA") != nullptr || + std::strstr(backend_name, "ROCm") != nullptr); + const bool direct_depthwise = transcribe::conformer::resolve_conv_direct( + "TRANSCRIBE_CONV_DIRECT_DW", "TRANSCRIBE_CONV_NO_DIRECT_DW", backend_direct); eb.n_blocks_local = (T_enc + hp.enc_context_size - 1) / hp.enc_context_size; const int T_pad = eb.n_blocks_local * hp.enc_context_size; eb.last_block_rem = T_enc - (eb.n_blocks_local - 1) * hp.enc_context_size; @@ -306,9 +352,9 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, transcribe::debug::mark_tensor_for_dump(x); } - ggml_tensor * conv_out = - conv_module(ctx, x, b, b.conv_bn_fused_scale, b.conv_bn_fused_bias, conv_k, static_cast(inner_dim)); - x = ggml_add(ctx, x, conv_out); + ggml_tensor * conv_out = conv_module(ctx, x, b, b.conv_bn_fused_scale, b.conv_bn_fused_bias, conv_k, + static_cast(inner_dim), direct_depthwise); + x = ggml_add(ctx, x, conv_out); if (i == 0) { named(x, "enc.block.0.post_conv"); eb.dumps.block_0_post_conv = x; diff --git a/src/arch/granite_nar/encoder.h b/src/arch/granite_nar/encoder.h index 1a24285f..2c2b7bd4 100644 --- a/src/arch/granite_nar/encoder.h +++ b/src/arch/granite_nar/encoder.h @@ -90,7 +90,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const GraniteNarWeights & weights, const GraniteNarHParams & hp, int T_enc, - bool use_flash); + bool use_flash, + const char * backend_name); // Shaw bookkeeping helpers (identical to AR granite). std::vector precompute_pos_rows(int context_size, int max_pos_emb); diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index e420483c..8390eeb9 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -562,7 +562,8 @@ transcribe_status run(transcribe_session * ctx_base, } // Encoder graph. - EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash); + EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash, + cm->backend.c_str()); if (eb.graph == nullptr || eb.cat_out == nullptr) { return TRANSCRIBE_ERR_GGUF; } From a1e85a6fc3e5df45bdf9fc106abaaa0a9c4be37c Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 17 Sep 2026 17:05:04 +0800 Subject: [PATCH 17/17] causal lm mem reduction --- src/arch/canary_qwen/model.cpp | 20 +---------------- src/arch/moss/decoder.cpp | 2 ++ src/arch/moss/decoder.h | 2 +- src/arch/moss/model.cpp | 13 +++-------- src/arch/moss/weights.h | 2 +- src/arch/qwen3_asr/model.cpp | 23 +------------------ src/arch/voxtral/decoder.cpp | 4 +++- src/arch/voxtral/decoder.h | 2 +- src/arch/voxtral/model.cpp | 26 +++------------------- src/arch/voxtral/weights.h | 4 +--- src/arch/voxtral_realtime/model.cpp | 13 +---------- src/causal_lm/causal_lm.cpp | 34 ++++++++++++++++------------- src/causal_lm/causal_lm.h | 6 +++-- 13 files changed, 41 insertions(+), 110 deletions(-) diff --git a/src/arch/canary_qwen/model.cpp b/src/arch/canary_qwen/model.cpp index 16f53fe4..b0a0b31e 100644 --- a/src/arch/canary_qwen/model.cpp +++ b/src/arch/canary_qwen/model.cpp @@ -713,24 +713,6 @@ transcribe_status init_context(transcribe_model * model, cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - auto * cm = static_cast(model); - { - ggml_type kv_type = GGML_TYPE_F16; - if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { - kv_type = GGML_TYPE_F32; - } - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, - /*n_ctx=*/2048, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, - cm->hparams.dec_n_layers, kv_type)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "canary_qwen init_context: KV cache allocation failed " - "(n_ctx=2048, %d kv-heads x %d head-dim x %d layers) — " - "out of memory.", - cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, cm->hparams.dec_n_layers); - return TRANSCRIBE_ERR_OOM; - } - } - *out_ctx = cc.release(); return TRANSCRIBE_OK; } @@ -924,7 +906,7 @@ transcribe_status run(transcribe_session * context, // hold prompt + generation budget, rounded up to a power of two (the step // graph's flash-attn path wants pow2 attention width). A pre-allocated // smaller cache is freed and re-allocated. - int want_n_ctx = 1024; + int want_n_ctx = 256; while (want_n_ctx < T_prompt + k_max_new) { want_n_ctx *= 2; } diff --git a/src/arch/moss/decoder.cpp b/src/arch/moss/decoder.cpp index 0623c177..98bcf916 100644 --- a/src/arch/moss/decoder.cpp +++ b/src/arch/moss/decoder.cpp @@ -36,6 +36,8 @@ causal_lm::BlockView to_block_view(const MossDecBlock & b) { v.attn_q_norm = b.attn_q_norm; v.attn_k_norm = b.attn_k_norm; v.ffn_gate_up_w = b.ffn_gate_up_w; + v.ffn_gate_w = b.ffn_gate_w; + v.ffn_up_w = b.ffn_up_w; v.ffn_down_w = b.ffn_down_w; return v; } diff --git a/src/arch/moss/decoder.h b/src/arch/moss/decoder.h index c812518b..c96cf85f 100644 --- a/src/arch/moss/decoder.h +++ b/src/arch/moss/decoder.h @@ -2,7 +2,7 @@ // // 28-layer Qwen3 causal decoder, identical block math to arch/qwen3_asr via // src/causal_lm/ (pre-LN RMSNorm, GQA with per-head q/k RMSNorm, NeoX RoPE -// theta 1e6, SwiGLU on packed gate+up, tied lm_head). +// theta 1e6, SwiGLU, tied lm_head). // // Audio injection differs from qwen3_asr: MOSS audio-pad positions are NOT // contiguous (the processor interleaves time-marker digit tokens into the diff --git a/src/arch/moss/model.cpp b/src/arch/moss/model.cpp index 3cacf0ff..af189b95 100644 --- a/src/arch/moss/model.cpp +++ b/src/arch/moss/model.cpp @@ -330,7 +330,9 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par } gguf_free(gguf_data); - { + // Keep the established one-matmul SwiGLU path on accelerators. On CPU, + // use the original gate/up projections to avoid retaining a packed copy. + if (m->plan.primary_kind != transcribe::BackendKind::Cpu) { std::vector entries; entries.reserve(m->weights.dec_blocks.size()); for (auto & b : m->weights.dec_blocks) { @@ -365,15 +367,6 @@ transcribe_status init_context(transcribe_model * model, cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - auto * cm = static_cast(model); - { - ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, /*n_ctx=*/2048, cm->hparams.dec_n_kv_heads, - cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, kv_type)) { - log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moss init_context: KV cache allocation failed"); - return TRANSCRIBE_ERR_OOM; - } - } *out_ctx = cc.release(); return TRANSCRIBE_OK; } diff --git a/src/arch/moss/weights.h b/src/arch/moss/weights.h index fcfd540b..a4d6c75b 100644 --- a/src/arch/moss/weights.h +++ b/src/arch/moss/weights.h @@ -159,7 +159,7 @@ struct MossDecBlock { ggml_tensor * ffn_gate_w = nullptr; ggml_tensor * ffn_up_w = nullptr; ggml_tensor * ffn_down_w = nullptr; - ggml_tensor * ffn_gate_up_w = nullptr; // packed at load + ggml_tensor * ffn_gate_up_w = nullptr; // packed on accelerator backends }; struct MossDecFinal { diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 4965f02e..3b87396e 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -300,27 +300,6 @@ transcribe_status init_context(transcribe_model * model, cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - // Pre-allocate KV cache at context creation so the first run - // doesn't pay the allocation cost inside the decode phase. - auto * cm = static_cast(model); - { - ggml_type kv_type = GGML_TYPE_F16; - if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { - kv_type = GGML_TYPE_F32; - } - const int initial_n_ctx = std::min(1024, qwen3_context_ceiling(cc->n_ctx, cm->hparams)); - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, initial_n_ctx, cm->hparams.dec_n_kv_heads, - cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, kv_type)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "qwen3_asr init_context: KV cache allocation failed " - "(n_ctx=%d, %d kv-heads x %d head-dim x %d layers) — " - "out of memory.", - initial_n_ctx, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, - cm->hparams.dec_n_layers); - return TRANSCRIBE_ERR_OOM; - } - } - *out_ctx = cc.release(); return TRANSCRIBE_OK; } @@ -727,7 +706,7 @@ transcribe_status run(transcribe_session * session, // hold prompt + generation budget, rounded up to a power of two (the step // graph's flash-attn path wants pow2 attention width). A pre-allocated // smaller cache is freed and re-allocated. - int want_n_ctx = 1024; + int want_n_ctx = 256; while (want_n_ctx < T_prompt + k_max_new) { want_n_ctx *= 2; } diff --git a/src/arch/voxtral/decoder.cpp b/src/arch/voxtral/decoder.cpp index c08217a8..c5724b36 100644 --- a/src/arch/voxtral/decoder.cpp +++ b/src/arch/voxtral/decoder.cpp @@ -1,7 +1,7 @@ // arch/voxtral/decoder.cpp - Voxtral LM prefill/step graph builders. // // Reference: VoxtralForConditionalGeneration's inner LlamaForCausalLM. The -// per-block math (pre-LN RMSNorm, GQA, NEOX RoPE, SwiGLU on packed gate_up) is +// per-block math (pre-LN RMSNorm, GQA, NEOX RoPE, SwiGLU) is // the shared causal_lm module with null Q/K-norm slots (Llama has no per-head // Q/K norm). This file owns graph allocation, audio injection (3-way concat), // dump naming, and the UNTIED lm_head. @@ -40,6 +40,8 @@ causal_lm::BlockView to_block_view(const VoxtralDecBlock & b) { v.attn_q_norm = nullptr; // Llama: no Q-norm v.attn_k_norm = nullptr; // Llama: no K-norm v.ffn_gate_up_w = b.ffn_gate_up_w; + v.ffn_gate_w = b.ffn_gate_w; + v.ffn_up_w = b.ffn_up_w; v.ffn_down_w = b.ffn_down_w; return v; } diff --git a/src/arch/voxtral/decoder.h b/src/arch/voxtral/decoder.h index 094608ae..bbf4ccf8 100644 --- a/src/arch/voxtral/decoder.h +++ b/src/arch/voxtral/decoder.h @@ -4,7 +4,7 @@ // - pre-LN RMSNorm (eps 1e-5) // - GQA (32 Q heads, 8 KV heads, head_dim 128), NO per-head Q/K norm // - NEOX RoPE (rotate_half) at theta 1e8 -// - SwiGLU MLP: down(silu(gate(x)) * up(x)) via packed gate+up +// - SwiGLU MLP: down(silu(gate(x)) * up(x)) // - UNTIED lm_head (dec.output.weight, separate from token_embd) // // The per-block math is the shared causal_lm module; this file owns graph diff --git a/src/arch/voxtral/model.cpp b/src/arch/voxtral/model.cpp index e490742b..3a8a1e8c 100644 --- a/src/arch/voxtral/model.cpp +++ b/src/arch/voxtral/model.cpp @@ -511,8 +511,9 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par } gguf_free(gguf_data); - // Pack gate+up for one-mul_mat SwiGLU. - { + // Keep the established one-matmul SwiGLU path on accelerators. On CPU, + // use the original gate/up projections to avoid retaining a packed copy. + if (m->plan.primary_kind != transcribe::BackendKind::Cpu) { std::vector entries; entries.reserve(m->weights.dec_blocks.size()); for (auto & b : m->weights.dec_blocks) { @@ -549,27 +550,6 @@ transcribe_status init_context(transcribe_model * model, cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - auto * cm = static_cast(model); - { - ggml_type kv_type = GGML_TYPE_F16; - if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { - kv_type = GGML_TYPE_F32; - } - // Initial allocation only: run()/run_batch() grow the KV cache per - // utterance up to the decoder's trained context (131072 for Mini-3B). - // 4096 covers short clips (~5 min) without a realloc on the hot path. - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, - /*n_ctx=*/4096, cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, - cm->hparams.dec_n_layers, kv_type)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "voxtral init_context: KV cache allocation failed " - "(n_ctx=4096, %d kv-heads x %d head-dim x %d layers) — " - "out of memory.", - cm->hparams.dec_n_kv_heads, cm->hparams.dec_head_dim, cm->hparams.dec_n_layers); - return TRANSCRIBE_ERR_OOM; - } - } - *out_ctx = cc.release(); return TRANSCRIBE_OK; } diff --git a/src/arch/voxtral/weights.h b/src/arch/voxtral/weights.h index ca533d6b..80a549e1 100644 --- a/src/arch/voxtral/weights.h +++ b/src/arch/voxtral/weights.h @@ -170,9 +170,7 @@ struct VoxtralDecBlock { ggml_tensor * ffn_gate_w = nullptr; // [hidden, intermediate] ggml_tensor * ffn_up_w = nullptr; // [hidden, intermediate] ggml_tensor * ffn_down_w = nullptr; // [intermediate, hidden] - // Packed gate+up: [hidden, 2*intermediate]. Filled at load time by - // causal_lm::pack_gate_up(); the graph uses one mul_mat for both. - ggml_tensor * ffn_gate_up_w = nullptr; + ggml_tensor * ffn_gate_up_w = nullptr; // packed on accelerator backends }; struct VoxtralDecFinal { diff --git a/src/arch/voxtral_realtime/model.cpp b/src/arch/voxtral_realtime/model.cpp index f3d112a2..69f8f663 100644 --- a/src/arch/voxtral_realtime/model.cpp +++ b/src/arch/voxtral_realtime/model.cpp @@ -320,23 +320,12 @@ transcribe_status init_context(transcribe_model * model, // decoder KV ring is constant-memory; the wall is dec_max_position). cc->n_ctx = transcribe_session_params_n_ctx(params); - auto * cm = static_cast(model); - // Encoder + decoder flash attention ON by default on every backend. Flash is // the numerical source-of-truth for the encoder. Override TRANSCRIBE_NO_FLASH=1. cc->encoder_use_flash = true; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); - ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; - if (!transcribe::causal_lm::kv_init(cc->kv_cache, cm->plan.primary, /*n_ctx=*/2048, cm->hparams.dec_n_kv_heads, - cm->hparams.dec_head_dim, cm->hparams.dec_n_layers, kv_type)) { - transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, - "voxtral_realtime init_context: KV cache allocation failed — " - "out of memory."); - return TRANSCRIBE_ERR_OOM; - } - *out_ctx = cc.release(); return TRANSCRIBE_OK; } @@ -672,7 +661,7 @@ transcribe_status forward_buffer(Session * cc, // Grow KV to fit n_audio positions. if (cc->kv_cache.n_ctx < n_audio + 1) { const ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; - int want = 2048; + int want = 256; while (want < n_audio + 1) { want *= 2; } diff --git a/src/causal_lm/causal_lm.cpp b/src/causal_lm/causal_lm.cpp index f8ecb2da..51f0d98a 100644 --- a/src/causal_lm/causal_lm.cpp +++ b/src/causal_lm/causal_lm.cpp @@ -41,6 +41,15 @@ ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * return y; } +ggml_tensor * ffn_swiglu(ggml_context * ctx, const BlockView & view, ggml_tensor * x) { + if (view.ffn_gate_up_w != nullptr) { + return ggml_swiglu(ctx, mul_mat_f32acc(ctx, view.ffn_gate_up_w, x)); + } + ggml_tensor * gate = ggml_silu(ctx, mul_mat_f32acc(ctx, view.ffn_gate_w, x)); + ggml_tensor * up = mul_mat_f32acc(ctx, view.ffn_up_w, x); + return ggml_mul(ctx, gate, up); +} + } // namespace void KvCache::free() { @@ -343,9 +352,8 @@ ggml_tensor * block_prefill(ggml_context * ctx, if (view.ffn_scale != nullptr) { ff_norm = ggml_mul(ctx, ff_norm, view.ffn_scale); } - ggml_tensor * gate_up = mul_mat_f32acc(ctx, view.ffn_gate_up_w, ff_norm); - ggml_tensor * ff = ggml_swiglu(ctx, gate_up); - ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); + ggml_tensor * ff = ffn_swiglu(ctx, view, ff_norm); + ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); x = ggml_add(ctx, x, ff); return x; @@ -459,9 +467,8 @@ ggml_tensor * block_step(ggml_context * ctx, if (view.ffn_scale != nullptr) { ff_norm = ggml_mul(ctx, ff_norm, view.ffn_scale); } - ggml_tensor * gate_up = mul_mat_f32acc(ctx, view.ffn_gate_up_w, ff_norm); - ggml_tensor * ff = ggml_swiglu(ctx, gate_up); - ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); + ggml_tensor * ff = ffn_swiglu(ctx, view, ff_norm); + ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); x = ggml_add(ctx, x, ff); return x; @@ -573,9 +580,8 @@ ggml_tensor * block_step_n(ggml_context * ctx, if (view.ffn_scale != nullptr) { ff_norm = ggml_mul(ctx, ff_norm, view.ffn_scale); } - ggml_tensor * gate_up = mul_mat_f32acc(ctx, view.ffn_gate_up_w, ff_norm); - ggml_tensor * ff = ggml_swiglu(ctx, gate_up); - ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); + ggml_tensor * ff = ffn_swiglu(ctx, view, ff_norm); + ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); x = ggml_add(ctx, x, ff); return x; @@ -685,9 +691,8 @@ ggml_tensor * block_step_batched(ggml_context * ctx, if (view.ffn_scale != nullptr) { ff_norm = ggml_mul(ctx, ff_norm, view.ffn_scale); } - ggml_tensor * gate_up = mul_mat_f32acc(ctx, view.ffn_gate_up_w, ff_norm); - ggml_tensor * ff = ggml_swiglu(ctx, gate_up); - ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); + ggml_tensor * ff = ffn_swiglu(ctx, view, ff_norm); + ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); x = ggml_add(ctx, x, ff); return x; @@ -789,9 +794,8 @@ ggml_tensor * block_prefill_batched(ggml_context * ctx, if (view.ffn_scale != nullptr) { ff_norm = ggml_mul(ctx, ff_norm, view.ffn_scale); } - ggml_tensor * gate_up = mul_mat_f32acc(ctx, view.ffn_gate_up_w, ff_norm); - ggml_tensor * ff = ggml_swiglu(ctx, gate_up); - ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); + ggml_tensor * ff = ffn_swiglu(ctx, view, ff_norm); + ff = mul_mat_f32acc(ctx, view.ffn_down_w, ff); x = ggml_add(ctx, x, ff); return x; diff --git a/src/causal_lm/causal_lm.h b/src/causal_lm/causal_lm.h index 46999d84..0e2c5c09 100644 --- a/src/causal_lm/causal_lm.h +++ b/src/causal_lm/causal_lm.h @@ -40,9 +40,11 @@ struct BlockView { // Voxtral's Ministral backbone); helpers skip the norm when null. ggml_tensor * attn_q_norm = nullptr; // [head_dim] per-head Q-norm, or null ggml_tensor * attn_k_norm = nullptr; // [head_dim] per-head K-norm, or null - // Packed gate+up filled by pack_gate_up at load time; the graph runs - // one mul_mat + ggml_swiglu instead of two mul_mats + manual silu·mul. + // Use the packed projection when available; otherwise run the original + // gate and up projections separately to avoid retaining a packed copy. ggml_tensor * ffn_gate_up_w = nullptr; // [hidden, 2·intermediate] + 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 per-layer FFN-branch scale (ff_norm *= ffn_scale, broadcast // over the token axis). Null for standard callers; voxtral_realtime's