Fix decode budget scaling

docs/input-limits.md

@@ -73,23 +73,26 @@ need and do not have a length gate.
 | --- | --- | --- |
 | 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`. |
 
-These families wrap an LLM-style decoder whose context window
-(`audio_tokens + prompt + generation`) is the binding constraint. The number of
-tokens a clip consumes is a deterministic function of its sample count
-(`n_samples → mel frames → fixed subsampling → audio tokens`), so the library
-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.
+A clip's sample count deterministically fixes its decoder prefill size or
+encoder frame count, so the library checks the relevant bound before running
+and rejects over-length input immediately. The rejection goes through the log
+callback, not raw stderr.
 
 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
-status `TRANSCRIBE_ERR_OUTPUT_TRUNCATED` while keeping the partial transcript
-readable (exactly like an aborted run); `transcribe_was_truncated(session)` is
-also set, and a `WARN` is logged. A truncated transcript is never returned as
-`TRANSCRIBE_OK` — a caller cannot mistake it for complete — and the partial
-output is never discarded. In `transcribe_run_batch` this is a per-utterance
-status (the whole-batch call still returns `TRANSCRIBE_OK`).
+long enough to exhaust the remaining budget mid-decode. There, the run returns
+the hard status `TRANSCRIBE_ERR_OUTPUT_TRUNCATED` while keeping the partial
+transcript readable (exactly like an aborted run);
+`transcribe_was_truncated(session)` is also set, and a `WARN` is logged. A
+truncated transcript is never returned as `TRANSCRIBE_OK` — a caller cannot
+mistake it for complete — and the partial output is never discarded. In
+`transcribe_run_batch` this is a per-utterance status (the whole-batch call
+still returns `TRANSCRIBE_OK`).
+
+Autoregressive families scale the decode budget with audio length, capped by
+the remaining decoder context. Lowering `n_ctx` can therefore make truncation
+more likely. For `canary` and `cohere`, input and output have separate encoder
+and decoder limits; `max_audio_ms` reports the encoder limit, not a recommended
+chunk size.
 
 ### 3. Soft window — warn and proceed
 
@@ -159,7 +162,7 @@ with `TRANSCRIBE_ERR_INPUT_TOO_LONG` (one-shot and batch) or surfaced via
 
 | Situation | Status | Log | Result |
 | --- | --- | --- | --- |
-| Input within limit | `TRANSCRIBE_OK` | — | full transcript |
+| Input within limit and decode completes | `TRANSCRIBE_OK` | — | full transcript |
 | Over-length, hard-cap family | `TRANSCRIBE_ERR_INPUT_TOO_LONG` | `ERROR` via callback | no transcript (rejected before the decode) |
 | Generation ran long mid-decode | `TRANSCRIBE_ERR_OUTPUT_TRUNCATED` | `WARN` via callback | partial transcript readable; `transcribe_was_truncated() == true` |
 | Over-window, soft-window family | `TRANSCRIBE_OK` | `WARN` via callback | full transcript (accuracy may be degraded) |

docs/porting/families/cohere.md

@@ -15,6 +15,9 @@ native Transformers. C++ CPU validation passes locally.
   `[en, ar]`; config omits top-level `vocab_size` — the converter falls
   back to `head.num_classes`; upstream repo is gated)
 
+Upstream recommends segmenting audio into 35 s clips; the port enforces the
+encoder's larger architectural limit. See `docs/input-limits.md`.
+
 ## References
 
 - Canonical reference: native Hugging Face Transformers

src/arch/canary/model.cpp

@@ -14,6 +14,7 @@
 #include "transcribe-arch.h"
 #include "transcribe-batch-util.h"
 #include "transcribe-debug.h"
+#include "transcribe-decode-budget.h"
 #include "transcribe-env.h"
 #include "transcribe-flash-policy.h"
 #include "transcribe-load-common.h"
@@ -214,10 +215,13 @@ constexpr float kBnEps = 1e-5f;
 //   (a) INPUT — the encoder rel-pos table (enc_pos_emb_max_len, ~400 s).
 //       T_enc must stay within it or the runtime table aliases past the
 //       trained range; gated up front. Drives max_audio_ms.
-//   (b) DECODER self-KV (dec_max_position) + 512 max-new cap bound the
-//       OUTPUT length; an overrun is kept as a partial and flagged via
+//   (b) DECODER self-KV (dec_max_position) bounds the OUTPUT length; an
+//       overrun is kept as a partial and flagged via
 //       transcribe_was_truncated(), not rejected.
 
+// Generation reserve: floor under the per-run decode budget.
+constexpr int k_gen_reserve = 512;
+
 // Predicted encoder frame count T_enc for a given mel frame count. The
 // FastConformer pre-encode downsamples time via stride-2, kernel-3, pad-1
 // convs; each stage maps T_in -> floor((T_in-1)/2)+1. We fold that exact
@@ -388,10 +392,16 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par
     // effective_max_audio_ms to the encoder bound regardless of n_ctx; the
     // decoder self-KV (which n_ctx does lower) only bounds transcript length.
     if (m->hparams.dec_max_position > 0) {
-        m->limits.has_context_cap        = true;
-        m->limits.audio_from_caps        = true;
-        m->limits.model_max_ctx          = m->hparams.dec_max_position;
-        m->limits.gen_reserve            = 512;  // run()'s max-new-tokens cap
+        m->limits.has_context_cap = true;
+        m->limits.audio_from_caps = true;
+        m->limits.model_max_ctx   = m->hparams.dec_max_position;
+        m->limits.gen_reserve     = k_gen_reserve;
+        // Encoder rate, for the decode budget only: audio_from_caps pins
+        // effective_max_audio_ms to the encoder bound, so this moves no limit.
+        if (m->hparams.enc_subsampling_factor > 0 && m->hparams.fe_hop_length > 0 && m->hparams.fe_sample_rate > 0) {
+            m->limits.ms_per_audio_token = static_cast<double>(m->hparams.enc_subsampling_factor) *
+                                           m->hparams.fe_hop_length * 1000.0 / m->hparams.fe_sample_rate;
+        }
         // Whisper-style decoder self-KV: dec_d_model per layer, K and V, no GQA.
         m->limits.kv_elems_per_ctx_token = (int64_t) m->hparams.dec_d_model * m->hparams.dec_n_layers * 2;
     }
@@ -1089,8 +1099,10 @@ transcribe_status run(transcribe_session *          session,
 
         cc->clear_result();
 
-        const int eos_id     = cm->hparams.eos_token_id;
-        const int max_tokens = std::min(512, cc->kv_cache.n_ctx - prompt_len);
+        const int eos_id = cm->hparams.eos_token_id;
+        const int max_tokens =
+            transcribe::pick_decode_budget(transcribe::predict_transcript_tokens(T_enc, cm->limits.ms_per_audio_token),
+                                           k_gen_reserve, prompt_len, cc->kv_cache.n_ctx);
 
         int next_token = 0;
         if (prompt_skip_softmax && db.argmax_out != nullptr) {
@@ -1608,15 +1620,17 @@ transcribe_status run_batch(transcribe_session *          session,
     }
 
     // Batched KV cache.
-    const int max_new  = 512;
-    int       max_n_kv = 1024;
-    while (max_n_kv < prompt_len + max_new) {
-        max_n_kv *= 2;
-    }
     // Decoder self-KV ceiling: dec_max_position, optionally lowered (never
     // raised) by the caller's n_ctx knob. Default knob (0) leaves it at
     // dec_max_position, so in-spec batched decode is unchanged.
     const int n_ctx_cap = canary_context_ceiling(cc->n_ctx, hp);
+    const int max_new =
+        transcribe::pick_decode_budget(transcribe::predict_transcript_tokens(T_enc_max, cm->limits.ms_per_audio_token),
+                                       k_gen_reserve, prompt_len, n_ctx_cap);
+    int max_n_kv = 1024;
+    while (max_n_kv < prompt_len + max_new) {
+        max_n_kv *= 2;
+    }
     if (max_n_kv > n_ctx_cap) {
         max_n_kv = n_ctx_cap;
     }

src/arch/canary_qwen/model.cpp

@@ -28,6 +28,7 @@
 #include "transcribe-arch.h"
 #include "transcribe-batch-util.h"
 #include "transcribe-debug.h"
+#include "transcribe-decode-budget.h"
 #include "transcribe-env.h"
 #include "transcribe-flash-policy.h"
 #include "transcribe-load-common.h"
@@ -116,9 +117,8 @@ constexpr float      kBnEps              = 1e-5f;
 // TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that fills the generation budget
 // before end-of-stream is flagged via transcribe_was_truncated().
 
-// Per-run generation budget. Keep in sync with the single-utterance and
-// batched step loops below.
-constexpr int k_max_new = 256;
+// Generation reserve: what the input gate keeps free, and the decode-budget floor.
+constexpr int k_gen_reserve = 256;
 
 // Effective decoder context ceiling, in tokens: the model's trained maximum
 // (decoder.max_position_embeddings, e.g. 40960), optionally lowered — never
@@ -148,7 +148,7 @@ int64_t canary_qwen_max_audio_ms(const CanaryQwenHParams & hp) {
     // counts; ~14 for canary_qwen). Advisory headroom, generous enough to
     // cover small template drift.
     constexpr int k_prompt_overhead = 32;
-    const int     max_audio_tokens  = hp.dec_max_position - k_prompt_overhead - k_max_new;
+    const int     max_audio_tokens  = hp.dec_max_position - k_prompt_overhead - k_gen_reserve;
     if (max_audio_tokens <= 0) {
         return 0;
     }
@@ -554,7 +554,7 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par
         m->limits.has_context_cap    = true;
         m->limits.model_max_ctx      = m->hparams.dec_max_position;
         m->limits.prompt_overhead    = 32;
-        m->limits.gen_reserve        = k_max_new;
+        m->limits.gen_reserve        = k_gen_reserve;
         // audio_tokens ≈ mel_frames / subsampling_factor ;
         //   mel_frames = ms*sr/(hop*1000)
         m->limits.ms_per_audio_token = static_cast<double>(m->hparams.enc_subsampling_factor) *
@@ -910,22 +910,25 @@ transcribe_status run(transcribe_session *          context,
     // Input-length gate: audio + prompt + generation must fit the decoder
     // context window. Reject an over-length clip here, before prefill/decode.
     const int ceiling = canary_qwen_context_ceiling(cc->n_ctx, hp);
-    if (T_prompt + k_max_new > ceiling) {
+    if (T_prompt + k_gen_reserve > ceiling) {
         transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                             "canary_qwen run: input too long — %d audio + %d prompt tokens "
                             "leave no room for output within the %d-token context (need %d). "
                             "Shorten the audio (see transcribe_capabilities.max_audio_ms) or "
                             "split it into segments.",
-                            T_enc, prefix_len + suffix_len, ceiling, T_prompt + k_max_new);
+                            T_enc, prefix_len + suffix_len, ceiling, T_prompt + k_gen_reserve);
         return TRANSCRIBE_ERR_INPUT_TOO_LONG;
     }
 
+    const int max_new = transcribe::pick_decode_budget(
+        transcribe::predict_transcript_tokens(T_enc, cm->limits.ms_per_audio_token), k_gen_reserve, T_prompt, ceiling);
+
     // KV cache init (grow-to-fit, clamped to the context ceiling). Size to
-    // hold prompt + generation budget, rounded up to a power of two (the step
+    // hold prompt + decode 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;
-    while (want_n_ctx < T_prompt + k_max_new) {
+    while (want_n_ctx < T_prompt + max_new) {
         want_n_ctx *= 2;
     }
     if (want_n_ctx > ceiling) {
@@ -1073,7 +1076,6 @@ transcribe_status run(transcribe_session *          context,
 
     // Step loop.
     const int32_t eos_id   = hp.eos_token_id;
-    const int     max_new  = k_max_new;
     int           cur_past = T_prompt;
 
     int max_n_kv = 1024;
@@ -1448,13 +1450,13 @@ transcribe_status run_batch(transcribe_session *          session,
 
         // Input-length gate (same as single-shot run()); reject this utterance,
         // the rest of the batch still runs.
-        if (T_prompt[b] + k_max_new > ceiling) {
+        if (T_prompt[b] + k_gen_reserve > ceiling) {
             transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR,
                                 "canary_qwen run_batch: utterance %d input too long — %d audio "
                                 "+ %d prompt tokens leave no room for output within the "
                                 "%d-token context (need %d). Shorten the audio (see "
                                 "transcribe_capabilities.max_audio_ms) or split it.",
-                                b, T_enc[b], T_prompt[b] - T_enc[b], ceiling, T_prompt[b] + k_max_new);
+                                b, T_enc[b], T_prompt[b] - T_enc[b], ceiling, T_prompt[b] + k_gen_reserve);
             fail_status[b] = TRANSCRIBE_ERR_INPUT_TOO_LONG;
             continue;
         }
@@ -1476,9 +1478,11 @@ transcribe_status run_batch(transcribe_session *          session,
         }
         return TRANSCRIBE_OK;
     }
-    T_enc_max          = std::max(1, T_enc_max);
-    const int max_new  = k_max_new;
-    int       max_n_kv = 1024;
+    T_enc_max = std::max(1, T_enc_max);
+    const int max_new =
+        transcribe::pick_decode_budget(transcribe::predict_transcript_tokens(T_enc_max, cm->limits.ms_per_audio_token),
+                                       k_gen_reserve, max_T_prompt, ceiling);
+    int max_n_kv = 1024;
     while (max_n_kv < max_T_prompt + max_new) {
         max_n_kv *= 2;
     }

src/arch/cohere/model.cpp

@@ -15,6 +15,7 @@
 #include "transcribe-arch.h"
 #include "transcribe-batch-util.h"
 #include "transcribe-debug.h"
+#include "transcribe-decode-budget.h"
 #include "transcribe-env.h"
 #include "transcribe-flash-policy.h"
 #include "transcribe-load-common.h"
@@ -416,6 +417,9 @@ transcribe_status promote_conv_pw_to_f32_on_cpu(CohereModel & m) {
 
 constexpr const char k_default_variant[] = "cohere-asr";
 
+// Generation reserve: floor under the per-run decode budget.
+constexpr int k_gen_reserve = 512;
+
 // Forward declarations for the Arch trait below.
 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 **);
@@ -477,7 +481,7 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par
         // Fixed control-token preamble (see run()'s prompt_pieces). Audio is
         // in cross-KV, so there is no audio-token overhead here.
         m->limits.prompt_overhead        = 10;
-        m->limits.gen_reserve            = 512;  // max-new-tokens cap in run()
+        m->limits.gen_reserve            = k_gen_reserve;
         // ms-per-audio-token = subsampling_factor * hop_length * 1000 / sr.
         m->limits.ms_per_audio_token     = static_cast<double>(m->hparams.enc_subsampling_factor) *
                                            m->hparams.fe_hop_length * 1000.0 / m->hparams.fe_sample_rate;
@@ -1038,8 +1042,10 @@ transcribe_status run(transcribe_session *          session,
         // Load-time validation guarantees eos_token_id >= 0; no
         // fallback is needed here. See the tokenizer.eos_id() check
         // in cohere::load() at the top of this file.
-        const int eos_id     = cm->hparams.eos_token_id;
-        const int max_tokens = std::min(512, cc->kv_cache.n_ctx - prompt_len);
+        const int eos_id = cm->hparams.eos_token_id;
+        const int max_tokens =
+            transcribe::pick_decode_budget(transcribe::predict_transcript_tokens(T_enc, cm->limits.ms_per_audio_token),
+                                           k_gen_reserve, prompt_len, cc->kv_cache.n_ctx);
 
         // Pick the first generated token. Fast path reads a single
         // int32 argmax that the GPU computed; debug path reads the
@@ -1563,14 +1569,16 @@ transcribe_status run_batch(transcribe_session *          session,
     }
 
     // ----- Allocate batched KV cache -----
-    const int max_new  = std::min(512, /*budget*/ 4096);
-    int       max_n_kv = 1024;
-    while (max_n_kv < prompt_len + max_new) {
-        max_n_kv *= 2;
-    }
     // Honor the session context cap (same ceiling the single-shot path uses),
     // not the raw model max — so a lowered n_ctx bounds batch decoder KV too.
     const int n_ctx_cap = cohere_dec_ctx_ceiling(cc->n_ctx, hp);
+    const int max_new =
+        transcribe::pick_decode_budget(transcribe::predict_transcript_tokens(T_enc_max, cm->limits.ms_per_audio_token),
+                                       k_gen_reserve, prompt_len, n_ctx_cap);
+    int max_n_kv = 1024;
+    while (max_n_kv < prompt_len + max_new) {
+        max_n_kv *= 2;
+    }
     if (max_n_kv > n_ctx_cap) {
         max_n_kv = n_ctx_cap;
     }