diff --git a/examples/bench/batch_bench.cpp b/examples/bench/batch_bench.cpp index 681b2983..4c4812af 100644 --- a/examples/bench/batch_bench.cpp +++ b/examples/bench/batch_bench.cpp @@ -2,16 +2,14 @@ // // Loads a model + session ONCE, then times transcribe_run_batch over a // sweep of batch sizes, reporting per-utterance latency and wall time per -// batch. This is the committed replacement for the throwaway harness that -// produced reports/perf/l40s/parakeet-batch/*.json — it measures only the -// compute (model load is excluded), so per_utt_ms is directly comparable +// batch. Model load is excluded, so per_utt_ms is directly comparable // across batch sizes. // // The batch is built by repeating a single input clip N times, which // isolates the device-occupancy effect of batching from input-length // variance. (Real workloads vary; this is a controlled throughput probe, // not a WER measurement.) Output is a JSON array on stdout, one object per -// batch size, matching the historical schema so existing plots keep working. +// batch size. // // Usage: // transcribe-batch-bench -m model.gguf clip.wav \ diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index a869a07d..4186caa4 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1,23 +1,9 @@ // transcribe-cli - example CLI driver for transcribe.cpp // -// Pass 2 stub: this CLI exercises the public ABI end-to-end with the -// stub library. It loads a WAV file (must be 16 kHz mono float32 source -// or downmixable to mono), prints duration / sample count, optionally -// asks the library to load a model (which will return NOT_IMPLEMENTED -// in pass 2), and prints the resulting status string. -// -// Once the loader and decoder land, this same CLI will run real -// transcription. Argument parsing, log sink wiring, and result printing -// will grow in place rather than being rewritten. -// -// Usage: -// transcribe-cli [options] audio.wav -// Options: -// -m, --model PATH GGUF model file (optional in pass 2) -// -l, --language ISO BCP-47-ish language hint (e.g. en, de) -// -t, --translate set task to TRANSLATE -// -q, --quiet suppress library log output -// -h, --help show this help +// Loads a 16 kHz mono float32 WAV (or downmixable to mono), loads a GGUF +// model, and transcribes it through the public ABI. Supports single-file +// and batch modes, offline and streaming paths, and per-family knobs. +// Run with --help for the full option list. #include "transcribe.h" #include "transcribe/parakeet.h" @@ -555,18 +541,16 @@ int main(int argc, char ** argv) { transcribe_log_set(log_cb, nullptr); } - // ---- Batch mode: --batch reads a file list, one wav path per - // line. Loads the model ONCE and reuses the context across all - // files. Outputs one JSONL line per file to stdout when - // --batch-jsonl is set, otherwise the same human-readable format - // as single-file mode. ------------------------------------------------ + // Batch mode: --batch reads a file list, one wav path per line. Loads + // the model ONCE and reuses the context across all files. Outputs one + // JSONL line per file to stdout when --batch-jsonl is set, otherwise + // the same human-readable format as single-file mode. if (!args.batch_file.empty()) { if (args.model_path.empty()) { std::fprintf(stderr, "error: --batch requires --model\n"); return EXIT_FAILURE; } - // Read the file list. std::vector wav_paths; { std::ifstream fin(args.batch_file); @@ -577,7 +561,6 @@ int main(int argc, char ** argv) { } std::string line; while (std::getline(fin, line)) { - // Trim trailing whitespace / carriage return. while (!line.empty() && (line.back() == '\n' || line.back() == '\r' || line.back() == ' ' || line.back() == '\t')) @@ -597,7 +580,6 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "batch: %zu files\n", wav_paths.size()); } - // Load model once. struct transcribe_model_load_params mp; transcribe_model_load_params_init(&mp); mp.backend = args.backend; mp.gpu_device = args.gpu_device; @@ -610,7 +592,6 @@ int main(int argc, char ** argv) { return EXIT_FAILURE; } - // Init context once (reused across all files via run()). struct transcribe_session_params cp; transcribe_session_params_init(&cp); cp.n_threads = args.n_threads; cp.n_ctx = args.n_ctx; @@ -801,7 +782,6 @@ int main(int argc, char ** argv) { for (size_t i = 0; i < wav_paths.size(); ++i) { const std::string & wav = wav_paths[i]; - // Load wav. std::vector pcm; std::string wav_err; if (!transcribe_cli::load_wav_mono_16k(wav, pcm, wav_err)) { @@ -931,7 +911,7 @@ int main(int argc, char ** argv) { } std::fflush(stdout); } - } // end serial else + } if (!args.batch_jsonl) { std::fprintf(stderr, "batch: %d ok, %d failed out of %zu\n", @@ -943,7 +923,7 @@ int main(int argc, char ** argv) { return n_fail > 0 ? EXIT_FAILURE : EXIT_SUCCESS; } - // ---- Single-file mode (original path). ----------------------------- + // Single-file mode. std::vector pcm; std::string load_err; if (!transcribe_cli::load_wav_mono_16k(args.wav_path, pcm, load_err)) { diff --git a/examples/common/wav.h b/examples/common/wav.h index 416551a9..d472de67 100644 --- a/examples/common/wav.h +++ b/examples/common/wav.h @@ -3,7 +3,7 @@ // Loads a RIFF WAV file as mono float32 PCM at 16 kHz, which is the only // input format the v1 transcribe runtime accepts. The WAV loader does NOT // resample. Files at any other sample rate are rejected with a message -// pointing the user at sox/ffmpeg, per PLAN.md "Sample Rate Policy". +// pointing the user at sox/ffmpeg. // // This file lives under examples/common/ on purpose: dr_wav and the WAV // loader are not part of the runtime library. The core library never diff --git a/include/transcribe.h b/include/transcribe.h index 21d02130..6c9871e7 100644 --- a/include/transcribe.h +++ b/include/transcribe.h @@ -37,30 +37,20 @@ * single startup-time install only. * - Calling transcribe_log_set() after threads or models exist is * UNSUPPORTED in 0.x. The implementation is data-race-free (the two - * atomics ensure no torn reads), but the API does NOT guarantee any - * of the following under concurrent or mid-run reconfiguration: - * (a) pair-atomic publication of (callback, userdata) - an emitter - * may observe a callback from one generation paired with a - * userdata from another, especially if two threads call - * transcribe_log_set concurrently; - * (b) that the old callback or old userdata will not be invoked - * after transcribe_log_set returns - an in-flight emission on - * another thread may already be holding the previous pair on - * its stack and about to call through it. - * A caller that needs mid-run reconfiguration must provide its own - * external synchronization between transcribe_log_set callers AND - * keep every previous callback target and userdata alive until it - * can prove no thread can still emit with them. A future revision - * may add a properly synchronized reconfiguration entry point with - * stronger guarantees. + * atomics ensure no torn reads), but under concurrent or mid-run + * reconfiguration the API does NOT guarantee pair-atomic publication + * of (callback, userdata), nor that the previous callback/userdata + * will not still be invoked after transcribe_log_set returns (an + * in-flight emission on another thread may already hold the old pair). + * A caller that needs mid-run reconfiguration must externally + * synchronize transcribe_log_set callers AND keep every previous + * callback target and userdata alive until no thread can still emit + * with them. * * ABI stability (0.x): * - This library is pre-1.0. The on-disk ABI MAY break between 0.x minor * releases. Consumers should rebuild against matching headers and * should not assume any layout, enum value, or symbol set is frozen. - * The `struct_size` mechanism described below is forward-ABI scaffolding - * we will lean on for the post-1.0 stable surface; treat it as a - * compatibility hook in the making, not a stability promise today. * * Params (size-aware structs): * - Every caller-owned public struct crossing the ABI carries `struct_size` @@ -208,12 +198,7 @@ typedef enum { TRANSCRIBE_ERR_UNSUPPORTED_VARIANT = 6, TRANSCRIBE_ERR_OOM = 7, TRANSCRIBE_ERR_BACKEND = 8, - /* - * Reserved. v1 accepts only 16 kHz mono float32 PCM and the library - * does not link a resampler, so there is currently no caller-supplied - * sample rate to validate. This code is preserved at its assigned - * value for the future "caller declares input rate" entry point. - */ + /* Reserved; not currently returned. */ TRANSCRIBE_ERR_SAMPLE_RATE = 9, TRANSCRIBE_ERR_UNSUPPORTED_LANGUAGE = 10, TRANSCRIBE_ERR_UNSUPPORTED_TASK = 11, @@ -244,21 +229,9 @@ typedef enum { * transcribe_*_init function" specifically. */ TRANSCRIBE_ERR_BAD_STRUCT_SIZE = 14, - /* - * Reserved. Phase 2 ships warn-and-proceed semantics for - * transcribe_run_params::pnc: a non-DEFAULT value against a model that - * does not advertise transcribe_model_supports(model, TRANSCRIBE_FEATURE_PNC) - * emits a WARN-level log message and proceeds with the model's - * default behavior. This code is preserved at its assigned value - * for a future opt-in strict-advisory mode in which the dispatcher - * would return this code instead of warning. Not currently - * returned. - */ + /* Reserved; not currently returned. */ TRANSCRIBE_ERR_UNSUPPORTED_PNC = 15, - /* - * Reserved. Symmetric with TRANSCRIBE_ERR_UNSUPPORTED_PNC for - * transcribe_run_params::itn. Not currently returned. - */ + /* Reserved; not currently returned. */ TRANSCRIBE_ERR_UNSUPPORTED_ITN = 16, /* * Returned by transcribe_run / transcribe_run_batch when the input @@ -494,9 +467,9 @@ typedef enum { * Each family that exposes a runtime PNC toggle reads this field; families * that do not (whisper, parakeet, ...) ignore it. The dispatcher emits a * WARN-level log message when a non-DEFAULT value is set against a model - * whose capabilities advertise supports_pnc == false, then proceeds with - * the model's default behavior. Use transcribe_capabilities::supports_pnc - * to pre-check. + * for which transcribe_model_supports(model, TRANSCRIBE_FEATURE_PNC) + * returns false, then proceeds with the model's default behavior. Use + * that probe to pre-check. * * DEFAULT (0): the family's shipped default. Zero-init via * transcribe_run_params_init() gives this value. The family @@ -522,8 +495,9 @@ enum transcribe_pnc_mode { * transformation that renders numbers, dates, currencies, etc. in formal * form ("twenty twenty four" → "2024", "ten dollars" → "$10"). Each * family that exposes a runtime ITN toggle reads this field; families - * that do not ignore it. Non-DEFAULT values against - * supports_itn == false emit a WARN and proceed with default behavior. + * that do not ignore it. Non-DEFAULT values against a model for which + * transcribe_model_supports(model, TRANSCRIBE_FEATURE_ITN) returns false + * emit a WARN and proceed with default behavior. * * DEFAULT (0): family default. Zero-init gives this value. * OFF: explicit ITN off. Supporting families emit verbatim @@ -993,13 +967,17 @@ TRANSCRIBE_API void transcribe_session_params_init( * * pnc: punctuation+capitalization runtime toggle. See * transcribe_pnc_mode. DEFAULT is always safe. Non-DEFAULT - * values against models with supports_pnc == false emit a - * WARN and proceed with the model's default behavior. + * values against models for which + * transcribe_model_supports(model, TRANSCRIBE_FEATURE_PNC) is + * false emit a WARN and proceed with the model's default + * behavior. * * itn: inverse text normalization runtime toggle. See * transcribe_itn_mode. DEFAULT is always safe. Non-DEFAULT - * values against models with supports_itn == false emit a - * WARN and proceed with the model's default behavior. + * values against models for which + * transcribe_model_supports(model, TRANSCRIBE_FEATURE_ITN) is + * false emit a WARN and proceed with the model's default + * behavior. * * language: source language hint as a BCP-47-ish short code, or * NULL to autodetect (only if the model supports it). @@ -1110,10 +1088,6 @@ TRANSCRIBE_API void transcribe_run_params_init( * anything similar in the future) live behind the * transcribe_model_supports() probe, keyed by transcribe_feature. * A new feature is +1 enum value with no struct_size churn. - * - * Trajectory note: revisit this shape if the bool cluster in the - * struct grows past ~5 or if a plugin-like family ABI is introduced. - * Neither condition is in flight today. */ struct transcribe_capabilities { uint64_t struct_size; @@ -1368,8 +1342,7 @@ TRANSCRIBE_API const char * transcribe_model_meta_val_str( /* ----------------------------------------------------------------------- */ /* - * Load a GGUF model from disk. The _file suffix is preserved to leave - * room for transcribe_model_load_buffer() later without renaming. + * Load a GGUF model from disk. * * params may be NULL for all library defaults. To customize, initialize * a struct with transcribe_model_load_params_init() and set fields. A @@ -1445,11 +1418,10 @@ TRANSCRIBE_API transcribe_status transcribe_open( struct transcribe_session ** out_session); /* - * Deprecated alias for transcribe_session_free, kept for source - * compatibility with the prior split open/close shape. Both honor - * owns_model and free the owned model when present, so a caller may use - * either with any session. New code should call transcribe_session_free - * directly. Passing NULL is a no-op. + * Deprecated alias for transcribe_session_free. Both honor owns_model and + * free the owned model when present, so a caller may use either with any + * session. New code should call transcribe_session_free directly. Passing + * NULL is a no-op. */ TRANSCRIBE_API void transcribe_close(struct transcribe_session * session); @@ -2046,7 +2018,7 @@ TRANSCRIBE_API transcribe_status transcribe_stream_get_text( * run_params->language not in the * model's declared language list. * - * Failure semantics — three stages, split by whether the previous + * Failure semantics — three checkpoints, split by whether the previous * result snapshot survives: * * 1. Dispatcher preflight. Top-level argument checks run first: @@ -2066,9 +2038,9 @@ TRANSCRIBE_API transcribe_status transcribe_stream_get_text( * result snapshot fully preserved and no FAILED transition, so a * caller-side config typo cannot destroy the prior utterance's * transcript. A family that does not implement stream_validate - * defers these checks to stage 3. + * performs these checks in stream_begin. * - * 3. Post-clear stream_begin. Once stages 1-2 pass, the dispatcher + * 3. Post-clear stream_begin. Once the preflight checks pass, the dispatcher * clears the result snapshot, enters ACTIVE, and calls the * family's stream_begin hook. Any non-OK status from this point * (config a non-preflighting family only catches here, allocation @@ -2255,9 +2227,7 @@ TRANSCRIBE_API int transcribe_tokenize( * encode_ms time to run the encoder forward pass on the most recent * transcribe_run. * decode_ms time to run the decoder (predictor + joint + token - * search) on the most recent transcribe_run. Currently - * always 0 since the decoder isn't wired yet; phase 5 - * fills this in. + * search) on the most recent transcribe_run. * * Output struct: caller initializes via transcribe_timings_init() (zero- * fill); the library writes only the fields that fit within diff --git a/include/transcribe/whisper.h b/include/transcribe/whisper.h index 5b68adad..9ed37a36 100644 --- a/include/transcribe/whisper.h +++ b/include/transcribe/whisper.h @@ -10,8 +10,9 @@ * * Whisper exposes substantial real model-specific knobs (13 fields). * The PNC/ITN toggles that other families share via transcribe_run_params - * do not apply: whisper does not advertise supports_pnc / supports_itn, - * and a non-DEFAULT pnc/itn against a whisper model produces a WARN. + * do not apply: transcribe_model_supports(model, TRANSCRIBE_FEATURE_PNC) + * and (..., TRANSCRIBE_FEATURE_ITN) both return false for whisper, and a + * non-DEFAULT pnc/itn against a whisper model produces a WARN. * Probe via transcribe_model_accepts_ext_kind before pointing * transcribe_run_params::family at this struct. * diff --git a/src/arch/canary/canary.h b/src/arch/canary/canary.h index 13bae45d..3fd4dba8 100644 --- a/src/arch/canary/canary.h +++ b/src/arch/canary/canary.h @@ -1,21 +1,8 @@ // arch/canary/canary.h - Canary multitask AED model and context types. // -// This header is INTERNAL to src/arch/canary/. It defines the concrete -// classes that derive from transcribe_model / transcribe_session for -// the NVIDIA Canary family (FastConformer encoder + Transformer decoder -// with 4-slot or 5-slot multitask prompt). -// -// Modeled on src/arch/cohere/cohere.h. Differences from cohere: -// - encoder shape mirrors parakeet's FastConformer, but every linear -// (Q/K/V/out, both macaron FFs, attention-pos projection, conv -// pointwise pair) carries a bias term. Parakeet is bias-free; -// canary is NOT. See weights.cpp for the full bias catalog. -// - decoder is structurally identical to cohere's but tensor names -// differ (dec.layer.{i}.norm{1,2,3} / {self_attn,cross_attn,ffn}.{q,k,v,o,up,down}) -// - LM head is UNTIED (explicit dec.head.{weight,bias}) -// - 180m-flash variant has an encoder->decoder projection -// (enc_d_model=512 -> dec_d_model=1024); other variants share -// d_model and skip the projection +// NVIDIA Canary: FastConformer encoder (parakeet shape but with biases on +// every linear) + autoregressive Transformer decoder (untied LM head; +// 180m-flash adds an enc->dec projection) with a 4/5-slot multitask prompt. #pragma once @@ -149,12 +136,9 @@ struct CanarySession final : public transcribe_session { std::vector enc_host; - // Per-stage flash-attention controls. Same rationale as cohere: - // encoder rel-pos MHSA dk depends on enc_d_model/n_heads (96 for - // 180m-flash, 128 for the larger variants); decoder dk=128 - // uniformly. ggml's Metal flash kernel currently lacks a path for - // some encoder dk values, so we default encoder flash off on - // Metal. The decoder defaults flash on. + // Per-stage flash-attention controls. ggml's Metal flash kernel lacks a + // path for some encoder rel-pos MHSA dk values (96 for 180m-flash, 128 + // otherwise), so encoder flash defaults off on Metal; decoder defaults on. bool encoder_use_flash = true; bool decoder_use_flash = true; diff --git a/src/arch/canary/capabilities.cpp b/src/arch/canary/capabilities.cpp index 1f5a6fa3..b7ef0698 100644 --- a/src/arch/canary/capabilities.cpp +++ b/src/arch/canary/capabilities.cpp @@ -15,10 +15,8 @@ void apply_family_invariants(transcribe_model & model) { // directions are an optional GGUF contract in stt.translation.pairs. caps.supports_translate = true; - // V1 port: timestamps are explicitly experimental upstream and - // out of scope (per intake known_risks). Advertise NONE; the GGUF - // KV stt.capability.timestamps overrides this when a future port - // wires up the timestamp decoder path. + // Timestamps out of scope. Advertise NONE; the GGUF KV + // stt.capability.timestamps overrides this once the path is wired up. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; // Feature bits: cancellation is wired at the run level. Canary diff --git a/src/arch/canary/decoder.cpp b/src/arch/canary/decoder.cpp index dc1907ee..7de56080 100644 --- a/src/arch/canary/decoder.cpp +++ b/src/arch/canary/decoder.cpp @@ -1,12 +1,6 @@ // arch/canary/decoder.cpp - Canary autoregressive Transformer decoder. // -// Modeled on src/arch/cohere/decoder.cpp. Differences from cohere: -// - per-sublayer dump points: dec.layer.{i}.{self_attn,cross_attn,ffn}.out -// so the C++ matches the reference dumper's hooks at layers -// {0, n_layers/2, n_layers-1} -// - LM head is UNTIED: applies w.head.weight + w.head.bias rather -// than reusing the token embedding -// - tensor names follow canary's GGUF (norm1/2/3, q/k/v/o, ffn.up/down) +// Untied LM head; per-sublayer dump points at layers {0, n_layers/2, n_layers-1}. #include "decoder.h" @@ -287,9 +281,8 @@ ggml_tensor * mha_cross_cached( return o; } -// Layers we attach per-sublayer dump points to. Mirrors the reference -// dumper's _block_indices convention: {0, n_layers/2, n_layers-1}, with -// dedup if the set collapses (n_layers <= 2). +// Layers we attach per-sublayer dump points to: {0, n_layers/2, n_layers-1} +// (deduped when the set collapses, n_layers <= 2). bool is_dump_layer(int i, int n_layers) { if (n_layers <= 0) return false; if (i == 0 || i == n_layers - 1) return true; @@ -419,13 +412,8 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, const auto & blk = w.dec_blocks[i]; const bool dump_this = is_dump_layer(i, hp.dec_n_layers) && i < 64; - // norm1 + self-attention. Reference dump hooks first_sub_layer - // output BEFORE the residual add — see TransformerDecoderBlock - // forward_preln in NeMo, which does - // self_attn_output = first_sub_layer(LN(x), ...) - // self_attn_output += residual - // and the forward_hook captures self_attn_output BEFORE the +=. - // So we dump self_out (the sublayer's contribution), not x. + // norm1 + self-attention. NeMo's forward_hook captures the sublayer + // output BEFORE the residual add, so we dump self_out, not x. { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm1_w, blk.norm1_b); ggml_tensor * self_out = mha_self_cached( @@ -447,9 +435,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, x = ggml_add(ctx, x, self_out); } - // norm2 + cross-attention. Same pattern — second_sub_layer - // returns the raw cross-attn output before NeMo adds the - // residual. + // norm2 + cross-attention (dump before residual, as above). { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm2_w, blk.norm2_b); ggml_tensor * cross_out = mha_cross_cached( @@ -469,8 +455,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, x = ggml_add(ctx, x, cross_out); } - // norm3 + FFN. third_sub_layer (PositionWiseFF) returns the FFN - // output before NeMo adds the residual. + // norm3 + FFN (dump before residual, as above). { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm3_w, blk.norm3_b); ggml_tensor * ff = ggml_mul_mat(ctx, blk.ffn_up_w, x_norm); @@ -498,7 +483,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, x = named(x, "dec.out_before_head"); db.dumps.out_before_head = x; - // Untied LM head: head.weight is its own tensor, not the token embedding. + // Untied LM head. ggml_tensor * logits = ggml_mul_mat(ctx, w.head.weight, x); if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); @@ -594,7 +579,6 @@ StepBuild build_step_graph(ggml_context * ctx, for (int i = 0; i < hp.dec_n_layers; ++i) { const auto & blk = w.dec_blocks[i]; - // norm1 + self-attention with set_rows-based KV write. { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm1_w, blk.norm1_b); ggml_tensor * self_out = mha_self_step( @@ -609,8 +593,7 @@ StepBuild build_step_graph(ggml_context * ctx, x = ggml_add(ctx, x, self_out); } - // norm2 + cross-attention. Unchanged from prompt path: cross_k/v - // are pre-populated full T_enc, no n_past dependency. + // norm2 + cross-attention. cross_k/v are pre-populated full T_enc. { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm2_w, blk.norm2_b); ggml_tensor * cross_out = mha_cross_cached( @@ -623,7 +606,6 @@ StepBuild build_step_graph(ggml_context * ctx, x = ggml_add(ctx, x, cross_out); } - // norm3 + FFN. { ggml_tensor * x_norm = layer_norm(ctx, x, blk.norm3_w, blk.norm3_b); ggml_tensor * ff = ggml_mul_mat(ctx, blk.ffn_up_w, x_norm); @@ -869,7 +851,7 @@ StepBuildBatched build_step_graph_batched(ggml_context * ctx, } x = layer_norm(ctx, x, w.dec_final.norm_w, w.dec_final.norm_b); - ggml_tensor * logits = ggml_mul_mat(ctx, w.head.weight, x); // untied head + ggml_tensor * logits = ggml_mul_mat(ctx, w.head.weight, x); if (w.head.bias != nullptr) logits = ggml_add(ctx, logits, w.head.bias); sb.argmax_out = ggml_argmax(ctx, logits); // [B] diff --git a/src/arch/canary/decoder.h b/src/arch/canary/decoder.h index 32afa6a9..52880b68 100644 --- a/src/arch/canary/decoder.h +++ b/src/arch/canary/decoder.h @@ -1,12 +1,6 @@ // arch/canary/decoder.h - Canary autoregressive Transformer decoder graph builder. // -// Modeled on src/arch/cohere/decoder.h. Differences: -// - per-sublayer dumps (self_attn / cross_attn / ffn outputs) at -// layers {0, n_layers/2, n_layers-1} so the C++ matches the -// reference dumper's hooks -// - LM head is UNTIED: head.weight is read as a separate tensor -// and `mul_mat(head.weight, x) + head.bias` produces the logits -// - tensor naming follows canary's GGUF (norm1/2/3, q/k/v/o, ffn.up/down) +// Untied LM head; per-sublayer dumps at layers {0, n_layers/2, n_layers-1}. #pragma once diff --git a/src/arch/canary/encoder.cpp b/src/arch/canary/encoder.cpp index 4de4719f..62ca3c09 100644 --- a/src/arch/canary/encoder.cpp +++ b/src/arch/canary/encoder.cpp @@ -1,12 +1,7 @@ // arch/canary/encoder.cpp - Canary FastConformer encoder graph builder. // -// Glue over transcribe::conformer. Encoder shape mirrors parakeet's -// FastConformer, but every linear (Q/K/V/out, both macaron FFs, the -// attention-pos projection, conv pointwise pair) carries a bias. Each -// BlockView bias slot is threaded through from the loader — see the -// per-block wiring below. The optional encoder->decoder projection is -// canary-specific and only fires for variants where -// stt.canary.decoder.encoder_decoder_proj=true (180m-flash today). +// Glue over transcribe::conformer: threads each BlockView bias slot through +// from the loader, plus the optional 180m-flash encoder->decoder projection. #include "encoder.h" @@ -28,9 +23,8 @@ namespace { namespace conf = transcribe::conformer; -// Match parakeet: direct dw on every backend for the in-block site; -// im2col for pre-encode (preserves parakeet's historical behavior since -// canary's encoder geometry is identical to parakeet's at the dw site). +// Direct dw on every backend for the in-block site; im2col for pre-encode +// (canary's encoder geometry matches parakeet's at the dw site). bool detect_direct_dw_in_block(const char * /*backend*/) { const char * env = std::getenv("TRANSCRIBE_CONV_DIRECT_DW"); if (env != nullptr) return true; @@ -53,13 +47,11 @@ conf::PreEncodeView to_view(const CanaryPreEncode & pe) { conf::BlockView to_view(const CanaryBlock & b) { conf::BlockView v; - // FF1 — every canary variant ships biases on the FFN linears. v.norm_ff1_w = b.norm_ff1_w; v.norm_ff1_b = b.norm_ff1_b; v.ff1_lin1_w = b.ff1_lin1_w; v.ff1_lin1_b = b.ff1_lin1_b; v.ff1_lin2_w = b.ff1_lin2_w; v.ff1_lin2_b = b.ff1_lin2_b; - // Self-attention with relative position — Q/K/V/out all have biases. v.norm_attn_w = b.norm_attn_w; v.norm_attn_b = b.norm_attn_b; v.attn_q_w = b.attn_q_w; v.attn_q_b = b.attn_q_b; @@ -70,7 +62,6 @@ conf::BlockView to_view(const CanaryBlock & b) { v.attn_pos_u = b.attn_pos_u; v.attn_pos_v = b.attn_pos_v; - // Conv module (biases on pw1/dw/pw2 — present in canary GGUF). v.norm_conv_w = b.norm_conv_w; v.norm_conv_b = b.norm_conv_b; v.conv_pw1_w = b.conv_pw1_w; v.conv_pw1_b = b.conv_pw1_b; @@ -79,7 +70,6 @@ conf::BlockView to_view(const CanaryBlock & b) { v.conv_bn_fused_scale = b.conv_bn_fused_scale; v.conv_bn_fused_bias = b.conv_bn_fused_bias; - // FF2 — same bias pattern as FF1. v.norm_ff2_w = b.norm_ff2_w; v.norm_ff2_b = b.norm_ff2_b; v.ff2_lin1_w = b.ff2_lin1_w; v.ff2_lin1_b = b.ff2_lin1_b; @@ -168,8 +158,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, transcribe::debug::mark_tensor_for_dump(x); } - // Mid block (indexed n_layers/2 to match the reference dumper's - // _block_indices convention {0, n//2, n-1}). + // Dump blocks {0, n/2, n-1}. const int mid_idx = hp.enc_n_layers / 2; // Blocks 1..N-1. @@ -193,9 +182,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // Native encoder output (pre-projection). Reference dumper emits - // this as enc.native — for variants without a projection it is - // numerically equal to enc.final. + // Native (pre-projection) output; equals enc.final when there is no proj. eb.dumps.native_out = x; { ggml_tensor * named_native = conf::named(x, "enc.native"); diff --git a/src/arch/canary/encoder.h b/src/arch/canary/encoder.h index a90eec9e..a72834d2 100644 --- a/src/arch/canary/encoder.h +++ b/src/arch/canary/encoder.h @@ -1,12 +1,7 @@ // arch/canary/encoder.h - Canary FastConformer encoder graph builder. // -// Encoder shape mirrors parakeet's FastConformer, but every linear -// (Q/K/V/out, both macaron FFs, attention-pos projection, conv -// pointwise pair) carries a bias term — the parakeet shape is -// bias-FREE, canary's is NOT. See weights.cpp for the full bias -// catalog. Only the optional encoder->decoder projection is -// canary-specific (180m-flash only — every other variant has -// enc_d_model == dec_d_model and skips it). +// Parakeet-shape FastConformer but with biases on every linear; 180m-flash +// adds an optional encoder->decoder projection. #pragma once diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index a08750eb..82f6fa50 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -1,14 +1,8 @@ // arch/canary/model.cpp - Canary multitask AED family handler. // // Load / init_context / run lifecycle for the four canary variants -// (180m-flash, 1b-flash, 1b-v2, 1b). Encoder is FastConformer in the -// parakeet shape, but unlike parakeet every linear (Q/K/V/out, both -// macaron FFs, attention-pos projection, conv pointwise pair) carries -// a bias term — see weights.cpp for the full set. Decoder is an -// autoregressive Transformer (cohere-shape, but with canary-specific -// tensor names and an UNTIED LM head). The multitask prompt is a -// 4-slot or 5-slot positional sequence read from the GGUF's -// stt.canary.special.* / stt.canary.tokenizer.prompt_format KV. +// (180m-flash, 1b-flash, 1b-v2, 1b): FastConformer encoder + autoregressive +// Transformer decoder, driven by a 4/5-slot multitask prompt. #include "canary.h" @@ -220,43 +214,18 @@ namespace { constexpr float kBnEps = 1e-5f; -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Canary is the same -// shape as cohere ASR: a FastConformer encoder feeding an autoregressive -// Transformer decoder, with TWO distinct limits: -// -// (a) INPUT limit — the encoder's relative positional-encoding table -// (enc_pos_emb_max_len, 5000 encoder frames for every shipped -// checkpoint). Like cohere (and unlike parakeet's unbounded -// conformer), encode_positions() is generated at the subsampled -// sequence length T_enc, so T_enc must stay within the trained -// pos-emb span or the runtime-generated table aliases past the -// trained range. T_enc is a deterministic function of mel frames, -// so this is the binding INPUT bound and it is gated up front, -// before the encoder graph is built. For the shipped rate -// (5000 frames * subsampling 8 * hop 160 / 16 kHz) that is ~400 s. -// -// (b) DECODER self-KV — dec_max_position (512 on canary-1b, 1024 on the -// flash / v2 variants) bounds prompt + transcript tokens, and a -// 512 max-new-tokens cap bounds generation. These bound the -// *output*: a transcript that runs long enough to exhaust the budget -// is kept as a partial result and flagged via -// transcribe_was_truncated(), not rejected. The prompt is only -// ~5-9 tokens, so this is an output-length wall, not an audio-length -// one — a 60 s clip of dense speech can hit it while a 300 s clip of -// sparse speech does not. -// -// The encoder is the input limit, so transcribe_capabilities::max_audio_ms -// is derived from (a); (b) is surfaced as the truncation flag. -// --------------------------------------------------------------------------- +// Input-length contract (see docs/input-limits.md). Two limits: +// (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 +// transcribe_was_truncated(), not rejected. // Predicted encoder frame count T_enc for a given mel frame count. The -// FastConformer pre-encode downsamples time by enc_subsampling_factor via -// stride-2, kernel-3, pad-1 convs; each stage maps T_in -> floor((T_in-1)/2)+1. -// We fold that exact per-stage recurrence so the prediction matches the -// graph's T_enc (the net result is floor(T_mel / subsampling_factor)). This -// is a pure host-side count; it touches no graph math. Mirrors cohere's -// cohere_predict_t_enc. +// 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 +// per-stage recurrence so the prediction matches the graph's T_enc. int canary_predict_t_enc(int mel_n_frames, int subsampling_factor) { if (mel_n_frames <= 0 || subsampling_factor <= 0) { return 0; @@ -281,17 +250,11 @@ int canary_predict_t_enc(int mel_n_frames, int subsampling_factor) { return t; } -// Advisory transcribe_capabilities::max_audio_ms: the longest audio whose -// encoder frame count T_enc still fits the trained positional-encoding span -// enc_pos_emb_max_len. Inverts the rate: -// T_enc = mel_frames / subsampling_factor -// mel_frames = ms * sample_rate / (hop_length * 1000) -// => ms = T_enc * subsampling_factor * hop_length * 1000 / sr -// at T_enc == enc_pos_emb_max_len. Returns 0 ("unknown / unbounded") if any -// rate hparam is missing, so a misconfigured model is never advertised with -// a wrong finite number. The decoder self-KV / max-new cap bounds the -// OUTPUT, not the input, so it does not enter this figure (a long-but-fitting -// clip may still truncate its transcript — see transcribe_was_truncated). +// Advisory max_audio_ms: longest audio whose T_enc still fits +// enc_pos_emb_max_len, inverting the rate +// ms = T_enc * subsampling_factor * hop_length * 1000 / sr +// at T_enc == enc_pos_emb_max_len. Returns 0 (unknown) if any rate hparam +// is missing, so a misconfigured model is never advertised with a wrong number. int64_t canary_max_audio_ms(const CanaryHParams & hp) { if (hp.enc_pos_emb_max_len <= 0 || hp.enc_subsampling_factor <= 0 || hp.fe_hop_length <= 0 || hp.fe_sample_rate <= 0) { @@ -302,10 +265,8 @@ int64_t canary_max_audio_ms(const CanaryHParams & hp) { return mel_frames * hp.fe_hop_length * 1000 / hp.fe_sample_rate; } -// Effective decoder self-KV ceiling, in tokens: the model's trained -// dec_max_position, optionally lowered — never raised — by the caller's -// session n_ctx knob. Used to size / clamp the autoregressive self-KV cache -// and bound generation. Mirrors cohere's cohere_dec_ctx_ceiling. +// Effective decoder self-KV ceiling, in tokens: dec_max_position, optionally +// lowered — never raised — by the caller's session n_ctx knob. int canary_context_ceiling(int32_t n_ctx_knob, const CanaryHParams & hp) { int ceiling = hp.dec_max_position > 0 ? hp.dec_max_position : 1024; if (n_ctx_knob > 0 && n_ctx_knob < ceiling) { @@ -432,22 +393,14 @@ transcribe_status load( return st; } - // Publish the input-length ceiling now that the encoder positional- - // encoding span and frontend rate are known (apply_family_invariants - // ran before the hparams were read, so it could not set this). The - // encoder is the binding INPUT limit for canary — see canary_max_audio_ms - // above and docs/input-limits.md. + // Publish the input-length ceiling now that the encoder positional span + // and frontend rate are known (apply_family_invariants ran before the + // hparams were read). The encoder is the binding INPUT limit for canary. m->caps.max_audio_ms = canary_max_audio_ms(m->hparams); - // Basis for the session-level limits query (transcribe_session_get_limits). - // Canary's INPUT bound is the ENCODER positional table (enc_pos_emb_max_len, - // ~400 s — see canary_max_audio_ms), not the decoder context: the decoder - // emits text tokens, so audio never consumes decoder context, and - // dec_max_position only bounds transcript LENGTH (surfaced via - // TRANSCRIBE_ERR_OUTPUT_TRUNCATED). So audio_from_caps = true keeps - // effective_max_audio_ms pinned to the encoder bound regardless of n_ctx, - // while effective_n_ctx / max_kv_bytes still reflect the decoder self-KV - // (which n_ctx does lower). Same shape as cohere. See docs/input-limits.md. + // Basis for transcribe_session_get_limits. audio_from_caps = true pins + // 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; @@ -606,9 +559,8 @@ transcribe_status init_context( cc->model = model; cc->n_threads = params->n_threads; cc->kv_type = params->kv_type; - // Cache the caller's n_ctx knob. For canary this lowers the DECODER - // self-KV ceiling (dec_max_position); it does not affect the encoder - // input limit. See canary_context_ceiling and docs/input-limits.md. + // Cache the caller's n_ctx knob. For canary this lowers the decoder + // self-KV ceiling only; it does not affect the encoder input limit. cc->n_ctx = transcribe_session_params_n_ctx(params); auto * cm = static_cast(model); @@ -625,37 +577,19 @@ transcribe_status init_context( return TRANSCRIBE_OK; } -// Build the multitask prompt token sequence. -// -// canary (4-slot, used by canary-1b): -// <|startoftranscript|> <|src_lang|> <|task|> <|target_lang|> <|nopnc|or|pnc|> -// -// canary2 (5-slot, used by canary-1b-flash, canary-1b-v2, canary-180m-flash): -// <|startofcontext|> <|startoftranscript|> <|src_lang|> <|target_lang|> -// <|task|> <|pnc|or|nopnc|> <|notimestamp|or|timestamp|> -// -// The actual prompt format used by NeMo is determined by the model's -// CanaryTokenizer / CanaryBPETokenizer implementation. For v1 we -// hardcode pnc=yes and timestamps=no (the first port focuses on -// transcription correctness; pnc/timestamp toggles land in a follow-up). -// -// Each slot's value is selected from the GGUF's stt.canary.special.*_id -// catalog at load time. Returns the assembled prompt or an empty -// vector + error message if any required slot is missing for the -// detected variant. +// Build the multitask prompt token sequence. Each slot's value comes from the +// GGUF's stt.canary.special.*_id catalog; returns {} if a required slot is +// missing. itn / timestamp / diarize are hardwired off; only pnc is toggleable. + +// canary-1 (4-slot), task token explicit: +// <|startoftranscript|> <|src_lang|> <|task|> <|target_lang|> <|pnc|> +// <|task|> is <|translate|> when src != tgt (and available), else <|transcribe|>. std::vector build_prompt_canary(const CanaryHParams & hp, int src_lang_id, int tgt_lang_id, const char * task, bool pnc) { - // canary-1 5-slot prompt template (from - // nemo.collections.common.prompts.canary.CanaryPromptFormatter): - // <|startoftranscript|> <|source_lang|> <|task|> <|target_lang|> <|pnc|> - // where <|task|> is one of <|transcribe|> / <|translate|> (canary-1 - // ships explicit task tokens, unlike canary2 which infers from - // src/tgt). We pick <|translate|> when src != tgt (and translate_id - // is available), else <|transcribe|>. std::vector ids; ids.reserve(8); @@ -692,13 +626,10 @@ std::vector build_prompt_canary2(const CanaryModel & cm, const char * /*task*/, bool pnc) { - // canary2 prompt template (from nemo.collections.common.prompts.canary2): + // canary2 prompt template: // <|startofcontext|> [decodercontext] <|startoftranscript|> // <|emo:?|> <|src_lang|> <|tgt_lang|> <|pnc|> <|itn|> <|timestamp|> <|diarize|> - // - // For ASR with empty decoder context the realized sequence is 9 - // tokens. itn / timestamp / diarize stay hardwired to their off - // tokens — only pnc is user-toggleable (see transcribe_run_params::pnc). + // ASR with empty decoder context realizes 9 tokens. std::vector ids; ids.reserve(9); @@ -790,7 +721,7 @@ transcribe_status run( transcribe::debug::init(); - // ----- Mel front-end ------------------------------------------- + // Mel front-end. if (!cm->mel.has_value()) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "canary run: model has no MelFrontend"); return TRANSCRIBE_ERR_INVALID_ARG; @@ -809,16 +740,10 @@ transcribe_status run( } cc->t_mel_us = ggml_time_us() - t_mel_start; - // ----- Input-length gate (see docs/input-limits.md) ------------- - // The encoder's relative positional-encoding table is the binding - // INPUT limit: encode_positions() is generated at the subsampled - // sequence length T_enc, so T_enc must stay within the trained span - // enc_pos_emb_max_len or the runtime pos table aliases past the - // trained range (silent out-of-bounds before this gate existed). - // T_enc is a deterministic function of mel_n_frames, so reject an - // over-length clip here — before the encoder graph is even built — - // rather than paying for a compute pass that cannot fit. The - // rejection goes through the log callback, not raw stderr. + // Input-length gate: T_enc must stay within enc_pos_emb_max_len or the + // runtime pos table aliases past the trained range. T_enc is a + // deterministic function of mel_n_frames, so reject an over-length clip + // here, before building the encoder graph. if (cm->hparams.enc_pos_emb_max_len > 0) { const int t_enc_pred = canary_predict_t_enc( mel_n_frames, cm->hparams.enc_subsampling_factor); @@ -835,14 +760,14 @@ transcribe_status run( } } - // ----- Reset compute state -------------------------------------- + // Reset compute state. if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; } cc->encoder_out = nullptr; - // ----- Build encoder graph -------------------------------------- + // Build encoder graph. { ggml_init_params init_params {}; init_params.mem_size = 8 * 1024 * 1024; @@ -919,7 +844,6 @@ transcribe_status run( transcribe::debug::dump_tensor("enc.pos_emb", eb.pos_emb_in, "encoder.pos_emb"); } - // Thread count. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); @@ -972,7 +896,7 @@ transcribe_status run( ggml_backend_tensor_get(eb.out, cc->enc_host.data(), 0, cc->enc_host.size() * sizeof(float)); - // ----- Build multitask prompt ----------------------------------- + // Build multitask prompt. const char * lang = (params && params->language) ? params->language : "en"; const bool is_translate = (params != nullptr && params->task == TRANSCRIBE_TASK_TRANSLATE); @@ -1010,8 +934,9 @@ transcribe_status run( // model's shipped behavior (pnc=on; matches the upstream model card's // published WER numbers). OFF / ON override explicitly. Non-DEFAULT // requests have already passed the dispatcher's advisory-warn gate - // (which only warns when supports_pnc == false; canary advertises - // supports_pnc = true so no warn fires here). + // (which only warns when transcribe_model_supports(model, + // TRANSCRIBE_FEATURE_PNC) is false; canary sets TRANSCRIBE_FEATURE_PNC + // so the probe returns true and no warn fires here). bool pnc = true; if (params != nullptr) { switch (params->pnc) { @@ -1037,7 +962,7 @@ transcribe_status run( } const int prompt_len = static_cast(prompt_ids.size()); - // ----- Init KV cache -------------------------------------------- + // Init KV cache. { if (cc->kv_cache.buffer != nullptr && cc->kv_cache.T_enc != T_enc) { cc->kv_cache.free(); @@ -1099,7 +1024,7 @@ transcribe_status run( return nullptr; }; - // ----- Cross-attention KV pre-compute --------------------------- + // Cross-attention KV pre-compute. const int64_t t_dec_start = ggml_time_us(); { if (!new_compute_ctx(4 * 1024 * 1024)) { @@ -1138,7 +1063,7 @@ transcribe_status run( cc->kv_cache.cross_populated = true; } - // ----- Prompt pass + autoregressive decode --------------------- + // Prompt pass + autoregressive decode. { if (!new_compute_ctx(4 * 1024 * 1024)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -1305,35 +1230,21 @@ transcribe_status run( cc->has_result = true; }; - // Two decoder loop variants; pick by primary backend kind. - // - // GPU (Vulkan/Metal/CUDA/SYCL): build_step_graph — one static- - // topology graph for the whole utterance. KV writes go via - // ggml_set_rows at runtime kv_idx; flash-attn reads a fixed - // max_n_kv window with a runtime mask. Removes per-step - // graph_build + sched_alloc, which dominate dispatch overhead - // on GPUs (~2 ms/tok on Vulkan-Renoir for canary-1b). - // - // CPU (incl. accelerator host-memory backends): build_decoder_ - // graph_kv per step — n_kv grows with n_past, attention only - // reads the populated prefix. CPU has no dispatch overhead - // to amortize, so the static-graph bandwidth tax (reading - // max_n_kv KV slots when avg n_past is ~25) is a net loss - // on the deep-decoder variants. - // - // The prompt pass uses build_decoder_graph_kv on both paths - // (already executed above) — only the per-token loop branches. + // Two decoder loop variants, picked by primary backend kind: + // GPU: build_step_graph — one static-topology graph reused for the + // whole utterance, amortizing per-step graph_build + sched_alloc + // dispatch overhead. + // CPU: build_decoder_graph_kv per step — n_kv grows with n_past, so + // no dispatch overhead to amortize and no static-graph bandwidth tax. 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 (primary_is_gpu) { - // ---------- Static-graph step path (GPU) ---------- - // max_n_kv: pad to next power of two with a 1024 floor. - // Vulkan/Metal flash-attn dispatches faster on pow2 ne[1] - // (~30% on M4 Max per qwen3_asr); the slight bandwidth cost - // is amortized by removing per-step graph_build + sched_alloc. + // Static-graph step path (GPU). max_n_kv: pad to next power of two + // with a 1024 floor — Vulkan/Metal flash-attn dispatches ~30% + // faster on pow2 ne[1]. int max_n_kv = 1024; while (max_n_kv < prompt_len + max_tokens) max_n_kv *= 2; if (max_n_kv > cc->kv_cache.n_ctx) max_n_kv = cc->kv_cache.n_ctx; @@ -1413,9 +1324,8 @@ transcribe_status run( if (next_token != eos_id) generated_ids.push_back(next_token); } } else { - // ---------- Dynamic-graph step path (CPU) ---------- - // Reserve the worst-case step graph so per-step alloc_graph - // doesn't grow the scheduler's memory pool. + // Dynamic-graph step path (CPU). Reserve the worst-case step graph + // so per-step alloc_graph doesn't grow the scheduler's memory pool. if (new_compute_ctx(4 * 1024 * 1024)) { const int worst_n_past = cc->kv_cache.n_ctx - 1; DecoderBuild db_reserve = build_decoder_graph_kv( @@ -1477,16 +1387,10 @@ transcribe_status run( } } - // The decode stopped either at EOS (complete) or at the generation - // cap / KV ceiling (truncated). Both step paths above share the same - // exit condition: the loop runs while `next_token != eos_id`, so a - // post-loop `next_token != eos_id` means it hit max_tokens / KV-full - // / a compute break without ever emitting end-of-stream. Surface that - // via transcribe_was_truncated() and a WARN rather than handing back - // a silently shortened transcript. A clean EOS decode leaves the flag - // false. See docs/input-limits.md. (The abort paths return early with - // a partial result and intentionally do NOT set the flag — abort is a - // caller-initiated stop, not a length truncation.) + // A post-loop next_token != eos_id means the decode hit max_tokens / + // KV-full / a compute break without end-of-stream: flag truncation and + // WARN rather than silently shortening. (Abort paths return early and + // intentionally do NOT set the flag — abort is not a length truncation.) if (next_token != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -1500,18 +1404,15 @@ transcribe_status run( commit_result(); } - // The partial transcript is committed by commit_result() above; a truncated - // decode returns the hard OUTPUT_TRUNCATED status (the result stays - // readable, like an aborted run). See docs/input-limits.md. + // Partial transcript committed above; a truncated decode returns the hard + // OUTPUT_TRUNCATED status (the result stays readable, like an aborted run). return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } // =========================================================================== -// Offline batched decode (transcribe_run_batch). Mirrors src/arch/cohere. -// Encoder stays serial per utterance (compute-bound); mel parallel; the -// autoregressive decode (self + cross attention) is batched. See -// docs/batching-autoregressive-plan.md. +// Offline batched decode (transcribe_run_batch). Encoder stays serial per +// utterance (compute-bound); mel parallel; the autoregressive decode is batched. // =========================================================================== // Encoder for one utterance from a precomputed mel buffer → host [hidden, T_enc]. @@ -1644,7 +1545,7 @@ transcribe_status run_batch( const int hidden = hp.dec_d_model; const int n_layer = hp.dec_n_layers; - // ----- Shared multitask prompt (identical across the batch) ----- + // Shared multitask prompt (identical across the batch). const char * lang = (params && params->language) ? params->language : "en"; const bool is_translate = (params != nullptr && params->task == TRANSCRIBE_TASK_TRANSLATE); @@ -1667,7 +1568,7 @@ transcribe_status run_batch( if (prompt_ids.empty()) return TRANSCRIBE_ERR_INVALID_ARG; const int prompt_len = static_cast(prompt_ids.size()); - // ----- Pass 0: parallel mel ----- + // Pass 0: parallel mel. std::vector valid(n, 0); std::vector> mel_bufs(n); std::vector mel_nf(n, 0); @@ -1685,12 +1586,9 @@ transcribe_status run_batch( }); mel_us += ggml_time_us() - t_mel0; - // ----- Pass 1: serial per-utterance encoder ----- - // Per-utterance input-length gate (same binding limit as run(): the - // encoder positional-encoding span enc_pos_emb_max_len). An over-length - // utterance is marked invalid with INPUT_TOO_LONG so the rest of the - // batch still decodes, mirroring how an encode failure drops one row. - // See docs/input-limits.md. + // Pass 1: serial per-utterance encoder. Per-utterance input-length gate + // (same enc_pos_emb_max_len limit as run()); an over-length utterance is + // marked invalid with INPUT_TOO_LONG so the rest of the batch still decodes. std::vector reject_status(n, TRANSCRIBE_ERR_INVALID_ARG); std::vector> enc_hosts(n); std::vector T_enc(n, 0); @@ -1728,7 +1626,7 @@ transcribe_status run_batch( return TRANSCRIBE_OK; } - // ----- Batched KV cache ----- + // 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; @@ -1769,7 +1667,7 @@ transcribe_status run_batch( const int64_t t_dec0 = ggml_time_us(); - // ----- Batched cross-attention K/V ----- + // Batched cross-attention K/V. { if (!new_compute_ctx(8 * 1024 * 1024)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -1801,12 +1699,10 @@ transcribe_status run_batch( cc->kv_cache.cross_populated = true; } - // ----- Batched step graph with a GROWING self-attention window ----- - // Self-KV holds only the short prompt + transcript (audio is in cross-KV), - // so the static n_ctx read is mostly empty for short clips. Start the read - // window at 64 and double it as n_past advances (rebuild graph + widen - // mask, O(log) rebuilds; KV cache persists). Cache capacity (max_n_kv) - // unchanged. See docs/batching-autoregressive-plan.md. + // Batched step graph with a GROWING self-attention window. Self-KV holds + // only the short prompt + transcript (audio is in cross-KV), so start the + // read window at 64 and double it as n_past advances (rebuild graph + widen + // mask, O(log) rebuilds; KV cache persists, capacity max_n_kv unchanged). const ggml_fp16_t f16_zero = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t f16_ninf = ggml_fp32_to_fp16(-INFINITY); const int32_t eos_id = hp.eos_token_id; @@ -1852,7 +1748,7 @@ transcribe_status run_batch( } const int64_t dec_us = ggml_time_us() - t_dec0; - // ----- Capture (strip control tokens like serial commit_result) ----- + // Capture (strip control tokens like serial commit_result). const bool strip = (params == nullptr) ? true : !params->keep_special_tags; const int valid_count = std::max(1, static_cast( std::count(valid.begin(), valid.end(), char(1)))); diff --git a/src/arch/canary/weights.cpp b/src/arch/canary/weights.cpp index 7f044730..c8d7fac1 100644 --- a/src/arch/canary/weights.cpp +++ b/src/arch/canary/weights.cpp @@ -1,16 +1,7 @@ // arch/canary/weights.cpp - read_canary_hparams + build_canary_weights. // -// Pattern follows parakeet/cohere weights.cpp. Differences: -// - encoder shape mirrors parakeet's FastConformer, but every linear -// (Q/K/V/out, both macaron FFs, attention-pos projection, conv -// pointwise pair) carries a bias. Parakeet is bias-free; canary is -// NOT. See the encoder-block loader below for the full bias set. -// - decoder tensor names use canary's scheme: dec.layer.{i}.norm{1,2,3} -// and dec.layer.{i}.{self_attn,cross_attn,ffn}.{q,k,v,o,up,down} -// - LM head is UNTIED: read dec.head.weight + dec.head.bias as -// standalone tensors, NOT a tied view of the token embedding -// - 180m-flash uniquely has enc.proj.{weight,bias} for the -// encoder->decoder dimension projection +// FastConformer encoder (biases on every linear) + autoregressive decoder +// (untied dec.head.{weight,bias}); 180m-flash adds enc.proj.{weight,bias}. #include "weights.h" @@ -47,7 +38,6 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, return TRANSCRIBE_ERR_INVALID_ARG; } - // Encoder. if (auto st = read_required_u32_kv(gguf, "stt.canary.encoder.n_layers", kFamilyTag, hp.enc_n_layers); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.canary.encoder.d_model", kFamilyTag, hp.enc_d_model); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.canary.encoder.n_heads", kFamilyTag, hp.enc_n_heads); st != TRANSCRIBE_OK) return st; @@ -58,7 +48,6 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, if (auto st = read_required_u32_kv(gguf, "stt.canary.encoder.pos_emb_max_len", kFamilyTag, hp.enc_pos_emb_max_len); st != TRANSCRIBE_OK) return st; if (auto st = read_optional_bool_kv(gguf, "stt.canary.encoder.use_bias", kFamilyTag, false, hp.enc_use_bias); st != TRANSCRIBE_OK) return st; - // Decoder. if (auto st = read_required_u32_kv(gguf, "stt.canary.decoder.n_layers", kFamilyTag, hp.dec_n_layers); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.canary.decoder.d_model", kFamilyTag, hp.dec_d_model); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.canary.decoder.n_heads", kFamilyTag, hp.dec_n_heads); st != TRANSCRIBE_OK) return st; @@ -112,22 +101,13 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, if (auto st = opt_special("stt.canary.special.transcribe_id", hp.transcribe_id); st != TRANSCRIBE_OK) return st; if (auto st = opt_special("stt.canary.special.translate_id", hp.translate_id); st != TRANSCRIBE_OK) return st; - // Language list -> language token id table. - // - // Two converter shapes for the routing keys: - // - canary-1 / 1b-flash / 180m-flash (CanaryTokenizer aggregate): - // stt.canary.tokenizer.lang_codes is ["en","de","es","fr", - // "spl_tokens"] (the per-language SP routing keys; "spl_tokens" - // is the special-vocab routing entry with no language id). - // - canary-1b-v2 (CanaryBPETokenizer): stt.canary.tokenizer.lang_codes - // is just ["all"] because there is a single SP — the actual - // language codes live only in general.languages. - // - // Either way the per-language token ids are written under - // stt.canary.special.lang._id. We walk both sources (lang_codes - // and general.languages), dedupe, and keep only codes whose _id KV - // resolves. This matches the runtime language list advertised by - // read_languages_kv but lets the hparam reader stand alone. + // Language list -> language token id table. Two converter shapes for the + // routing keys: aggregate tokenizers put per-language codes in + // stt.canary.tokenizer.lang_codes (+ a "spl_tokens" entry with no id); + // single-SP (canary-1b-v2) puts ["all"] there and the real codes in + // general.languages. Either way the ids live under + // stt.canary.special.lang._id, so walk both sources, dedupe, and + // keep only codes whose _id KV resolves. { std::vector all_codes; std::vector tmp; @@ -175,7 +155,6 @@ transcribe_status read_canary_hparams(const gguf_context * gguf, } } - // Frontend. if (auto st = read_required_string_kv(gguf, "stt.frontend.type", kFamilyTag, hp.fe_type); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.num_mels", kFamilyTag, hp.fe_num_mels); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.sample_rate", kFamilyTag, hp.fe_sample_rate); st != TRANSCRIBE_OK) return st; @@ -328,7 +307,7 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, const int64_t pre_encode_freq = channels * (hp.fe_num_mels / hp.enc_subsampling_factor); const int64_t pre_encode_in = pre_encode_freq; - // ----- pre_encode ----- + // pre_encode GET_CONV(weights.pre_encode.conv0_w, "enc.pre_encode.conv.0.weight", 3, 3, 1, channels); GET_F32 (weights.pre_encode.conv0_b, "enc.pre_encode.conv.0.bias", channels); GET_CONV(weights.pre_encode.conv2_w, "enc.pre_encode.conv.2.weight", 3, 3, 1, channels); @@ -342,18 +321,13 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_LIN (weights.pre_encode.out_w, "enc.pre_encode.out.weight", pre_encode_in, d_enc); GET_F32 (weights.pre_encode.out_b, "enc.pre_encode.out.bias", d_enc); - // ----- encoder blocks ----- - // Every linear in the canary encoder block carries a bias term — - // both macaron FFs, Q/K/V/out, attention-pos projection, and the - // conv pointwise pair. This differs from parakeet (bias-free - // linears) and is load-bearing: a previous comment claimed - // "bias-free" and that mismatch hid a real loader drift; do NOT - // re-introduce that wording without checking the GGUF. + // Encoder blocks. Every linear carries a bias term — both macaron FFs, + // Q/K/V/out, attention-pos projection, and the conv pointwise pair + // (unlike parakeet, which is bias-free). weights.blocks.assign(hp.enc_n_layers, CanaryBlock{}); for (int i = 0; i < hp.enc_n_layers; ++i) { auto & b = weights.blocks[i]; - // Macaron FF1 — canary always has biases on the linears. GET_F32(b.norm_ff1_w, lname("enc.blocks.%d.norm_ff1.weight", i), d_enc); GET_F32(b.norm_ff1_b, lname("enc.blocks.%d.norm_ff1.bias", i), d_enc); GET_LIN(b.ff1_lin1_w, lname("enc.blocks.%d.ff1.linear1.weight", i), d_enc, d_ff_e); @@ -361,7 +335,6 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_LIN(b.ff1_lin2_w, lname("enc.blocks.%d.ff1.linear2.weight", i), d_ff_e, d_enc); GET_F32(b.ff1_lin2_b, lname("enc.blocks.%d.ff1.linear2.bias", i), d_enc); - // Self-attention with relative position — Q/K/V/out also have biases. GET_F32(b.norm_attn_w, lname("enc.blocks.%d.norm_attn.weight", i), d_enc); GET_F32(b.norm_attn_b, lname("enc.blocks.%d.norm_attn.bias", i), d_enc); GET_LIN(b.attn_q_w, lname("enc.blocks.%d.attn.linear_q.weight", i), d_enc, d_enc); @@ -376,7 +349,6 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_F32(b.attn_pos_u, lname("enc.blocks.%d.attn.pos_bias_u", i), head_dim, n_heads); GET_F32(b.attn_pos_v, lname("enc.blocks.%d.attn.pos_bias_v", i), head_dim, n_heads); - // Conv module — pointwise/depthwise carry biases per the GGUF. GET_F32 (b.norm_conv_w, lname("enc.blocks.%d.norm_conv.weight", i), d_enc); GET_F32 (b.norm_conv_b, lname("enc.blocks.%d.norm_conv.bias", i), d_enc); GET_CONV(b.conv_pw1_w, lname("enc.blocks.%d.conv.pointwise1.weight", i), 1, d_enc, 2 * d_enc); @@ -390,7 +362,6 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_F32 (b.conv_bn_rm, lname("enc.blocks.%d.conv.bn.running_mean", i), d_enc); GET_F32 (b.conv_bn_rv, lname("enc.blocks.%d.conv.bn.running_var", i), d_enc); - // Macaron FF2 — biases present (same as FF1). GET_F32(b.norm_ff2_w, lname("enc.blocks.%d.norm_ff2.weight", i), d_enc); GET_F32(b.norm_ff2_b, lname("enc.blocks.%d.norm_ff2.bias", i), d_enc); GET_LIN(b.ff2_lin1_w, lname("enc.blocks.%d.ff2.linear1.weight", i), d_enc, d_ff_e); @@ -398,18 +369,17 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_LIN(b.ff2_lin2_w, lname("enc.blocks.%d.ff2.linear2.weight", i), d_ff_e, d_enc); GET_F32(b.ff2_lin2_b, lname("enc.blocks.%d.ff2.linear2.bias", i), d_enc); - // Final per-block layer norm. GET_F32(b.norm_out_w, lname("enc.blocks.%d.norm_out.weight", i), d_enc); GET_F32(b.norm_out_b, lname("enc.blocks.%d.norm_out.bias", i), d_enc); } - // ----- optional encoder->decoder projection (180m-flash only) ----- + // Optional encoder->decoder projection (180m-flash only). if (hp.dec_has_encoder_decoder_proj) { GET_LIN(weights.enc_proj.weight, "enc.proj.weight", d_enc, d_dec); GET_F32(weights.enc_proj.bias, "enc.proj.bias", d_dec); } - // ----- decoder embedding ----- + // Decoder embedding. { ggml_tensor * tw = ggml_get_tensor(ctx_meta, "dec.embed.token.weight"); if (tw == nullptr) { @@ -429,12 +399,11 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_F32(weights.dec_embed.norm_w, "dec.embed.norm.weight", d_dec); GET_F32(weights.dec_embed.norm_b, "dec.embed.norm.bias", d_dec); - // ----- decoder blocks (canary tensor naming) ----- + // Decoder blocks (canary tensor naming). weights.dec_blocks.assign(hp.dec_n_layers, CanaryDecBlock{}); for (int i = 0; i < hp.dec_n_layers; ++i) { auto & db = weights.dec_blocks[i]; - // norm1 + self-attention. GET_F32(db.norm1_w, lname("dec.layer.%d.norm1.weight", i), d_dec); GET_F32(db.norm1_b, lname("dec.layer.%d.norm1.bias", i), d_dec); GET_LIN(db.self_q_w, lname("dec.layer.%d.self_attn.q.weight", i), d_dec, d_dec); @@ -446,7 +415,6 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_LIN(db.self_o_w, lname("dec.layer.%d.self_attn.o.weight", i), d_dec, d_dec); GET_F32(db.self_o_b, lname("dec.layer.%d.self_attn.o.bias", i), d_dec); - // norm2 + cross-attention. GET_F32(db.norm2_w, lname("dec.layer.%d.norm2.weight", i), d_dec); GET_F32(db.norm2_b, lname("dec.layer.%d.norm2.bias", i), d_dec); GET_LIN(db.cross_q_w, lname("dec.layer.%d.cross_attn.q.weight", i), d_dec, d_dec); @@ -458,7 +426,6 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_LIN(db.cross_o_w, lname("dec.layer.%d.cross_attn.o.weight", i), d_dec, d_dec); GET_F32(db.cross_o_b, lname("dec.layer.%d.cross_attn.o.bias", i), d_dec); - // norm3 + FFN. GET_F32(db.norm3_w, lname("dec.layer.%d.norm3.weight", i), d_dec); GET_F32(db.norm3_b, lname("dec.layer.%d.norm3.bias", i), d_dec); GET_LIN(db.ffn_up_w, lname("dec.layer.%d.ffn.up.weight", i), d_dec, d_ff_d); @@ -467,11 +434,11 @@ transcribe_status build_canary_weights(ggml_context * ctx_meta, GET_F32(db.ffn_down_b, lname("dec.layer.%d.ffn.down.bias", i), d_dec); } - // ----- decoder final norm ----- + // Decoder final norm. GET_F32(weights.dec_final.norm_w, "dec.norm.weight", d_dec); GET_F32(weights.dec_final.norm_b, "dec.norm.bias", d_dec); - // ----- LM head (untied) ----- + // LM head (untied). GET_LIN(weights.head.weight, "dec.head.weight", d_dec, vocab); GET_F32(weights.head.bias, "dec.head.bias", vocab); diff --git a/src/arch/canary/weights.h b/src/arch/canary/weights.h index c74c782b..f8bfa197 100644 --- a/src/arch/canary/weights.h +++ b/src/arch/canary/weights.h @@ -1,15 +1,8 @@ // arch/canary/weights.h - Canary tensor catalog and per-instance weight slots. // -// Canary's encoder is a FastConformer in the parakeet shape but with -// biases on every linear (Q/K/V/out, both macaron FFs, attention-pos -// projection, conv pointwise pair) — unlike parakeet, which is -// bias-free. Its decoder is an autoregressive Transformer (cohere -// shape, but with distinct tensor names and an UNTIED LM head). -// -// 180m-flash uniquely has an encoder->decoder projection between the -// final encoder block and the decoder cross-attention K/V source -// (enc_d_model=512 -> dec_d_model=1024). The other three variants -// share enc_d_model=dec_d_model and skip the projection. +// FastConformer encoder (parakeet shape, biases on every linear) + +// autoregressive Transformer decoder (untied LM head). 180m-flash adds an +// enc->dec projection; other variants share enc_d_model=dec_d_model. #pragma once @@ -138,13 +131,9 @@ struct CanaryPreEncode { ggml_tensor * out_b = nullptr; }; -// One FastConformer encoder block. Despite the `stt.canary.encoder.use_bias` -// KV reading false in some GGUFs (a converter mislabel — the .bias tensors -// are always emitted), every canary variant ships biases on the FFN and -// attention linears. The cohere-style structure (with biases everywhere) -// is the right shape for canary. +// One FastConformer encoder block. Every canary variant ships biases on the +// FFN and attention linears, regardless of the stt.canary.encoder.use_bias KV. struct CanaryBlock { - // Macaron FF1. ggml_tensor * norm_ff1_w = nullptr; ggml_tensor * norm_ff1_b = nullptr; ggml_tensor * ff1_lin1_w = nullptr; @@ -152,7 +141,6 @@ struct CanaryBlock { ggml_tensor * ff1_lin2_w = nullptr; ggml_tensor * ff1_lin2_b = nullptr; - // Self-attention with relative positional encoding. ggml_tensor * norm_attn_w = nullptr; ggml_tensor * norm_attn_b = nullptr; ggml_tensor * attn_q_w = nullptr; @@ -167,7 +155,7 @@ struct CanaryBlock { ggml_tensor * attn_pos_u = nullptr; ggml_tensor * attn_pos_v = nullptr; - // Convolution module (BN tensors fused at load time into _scale/_bias). + // Conv module — BN fused at load into _scale/_bias. ggml_tensor * norm_conv_w = nullptr; ggml_tensor * norm_conv_b = nullptr; ggml_tensor * conv_pw1_w = nullptr; @@ -183,7 +171,6 @@ struct CanaryBlock { ggml_tensor * conv_bn_fused_scale = nullptr; ggml_tensor * conv_bn_fused_bias = nullptr; - // Macaron FF2. ggml_tensor * norm_ff2_w = nullptr; ggml_tensor * norm_ff2_b = nullptr; ggml_tensor * ff2_lin1_w = nullptr; @@ -191,7 +178,6 @@ struct CanaryBlock { ggml_tensor * ff2_lin2_w = nullptr; ggml_tensor * ff2_lin2_b = nullptr; - // Final per-block layer norm. ggml_tensor * norm_out_w = nullptr; ggml_tensor * norm_out_b = nullptr; }; @@ -214,7 +200,6 @@ struct CanaryDecEmbed { // norm2 -> cross_attn -> add residual // norm3 -> ffn -> add residual struct CanaryDecBlock { - // norm1 + self-attention. ggml_tensor * norm1_w = nullptr; ggml_tensor * norm1_b = nullptr; ggml_tensor * self_q_w = nullptr; @@ -226,7 +211,6 @@ struct CanaryDecBlock { ggml_tensor * self_o_w = nullptr; ggml_tensor * self_o_b = nullptr; - // norm2 + cross-attention. ggml_tensor * norm2_w = nullptr; ggml_tensor * norm2_b = nullptr; ggml_tensor * cross_q_w = nullptr; @@ -238,7 +222,6 @@ struct CanaryDecBlock { ggml_tensor * cross_o_w = nullptr; ggml_tensor * cross_o_b = nullptr; - // norm3 + FFN. ggml_tensor * norm3_w = nullptr; ggml_tensor * norm3_b = nullptr; ggml_tensor * ffn_up_w = nullptr; diff --git a/src/arch/canary_qwen/canary_qwen.h b/src/arch/canary_qwen/canary_qwen.h index b34eb1af..8bdffbac 100644 --- a/src/arch/canary_qwen/canary_qwen.h +++ b/src/arch/canary_qwen/canary_qwen.h @@ -61,28 +61,20 @@ struct CanaryQwenModel final : public transcribe_model { transcribe::BackendPlan plan; ggml_backend_buffer_t backend_buffer = nullptr; - // BatchNorm fusion in the conformer conv module — same as canary. - // Each block stores fused (scale, bias) computed at load time; - // running_mean/running_var are NOT used at inference. + // BatchNorm fusion in the conformer conv module: fused (scale, bias) + // computed at load time; running_mean/var unused at inference. ggml_context * bn_fused_ctx = nullptr; ggml_backend_buffer_t bn_fused_buffer = nullptr; - // F16/BF16 conv pointwise weights promoted to F32 on CPU backend so - // ggml's CPU kernels hit the F32 matmul path (avoids unhealthy - // fp16/bf16 round-tripping inside the inner loop). + // F16 conv pointwise/depthwise weights promoted to F32 (see model.cpp). ggml_context * conv_pw_f32_ctx = nullptr; ggml_backend_buffer_t conv_pw_f32_buffer = nullptr; - // BF16 encoder + decoder linear weights promoted to F32 on CPU - // backend. canary-1b-flash (the source of these encoder weights) - // shipped F32 GGUF and was validated F32-only; ggml's CPU BF16 - // matmul path was never exercised. Promotion is conservative and - // matches the reference's F32 inference regime. + // BF16 encoder + decoder linear weights promoted to F32 on CPU (see + // model.cpp; matches the reference's F32 inference regime). ggml_context * linear_f32_ctx = nullptr; ggml_backend_buffer_t linear_f32_buffer = nullptr; - // Qwen3 packed gate+up MLP weights — same as funasr_nano / - // qwen3_asr. transcribe::causal_lm::PackedGateUpHandles packed_gate_up; // C++ mel frontend (constructed once at load). diff --git a/src/arch/canary_qwen/decoder.cpp b/src/arch/canary_qwen/decoder.cpp index 79ca07e6..b5af743f 100644 --- a/src/arch/canary_qwen/decoder.cpp +++ b/src/arch/canary_qwen/decoder.cpp @@ -1,15 +1,9 @@ // arch/canary_qwen/decoder.cpp - Qwen3-1.7B LM prefill / step graphs. // -// Same Qwen3 block math as funasr_nano (and qwen3_asr) — verified -// byte-for-byte against HF's `transformers/models/qwen3/modeling_qwen3.py` -// (qwen_asr's `Qwen3ASRTextDecoderLayer` is HF's `Qwen3DecoderLayer` -// with only the class name changed). Funasr_nano validates the same -// helpers on Qwen3-0.6B with `dec.logits_raw.prefill` drift of -// max_abs=0.12 / mean_abs=0.022 (fp32 noise). -// -// Audio injection: 3-way concat. For B=1 with one -// `<|audioplaceholder|>` in the prompt, this is mathematically -// equivalent to NeMo SALM's `replace_placeholders_and_build_targets`. +// Same Qwen3 block math as funasr_nano and qwen3_asr (shared via +// src/causal_lm/). Audio injection: 3-way concat, equivalent to NeMo SALM's +// `replace_placeholders_and_build_targets` for B=1 with one +// `<|audioplaceholder|>`. #include "decoder.h" @@ -162,15 +156,8 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, x_suffix = ggml_cont(ctx, x_suffix); } - // Note: SALM's reference dumper captures `dec.token_emb` from the - // top-level embed_tokens call, which embeds the *un-expanded* text - // prompt (15 tokens, including a single audio_locator placeholder). - // Our C++ pipeline uses already-expanded input_ids (152 tokens, with - // T_audio locator slots that get overwritten by audio frames) and - // does not have a clean intermediate that matches the reference's - // (15, hidden) shape. We skip the C++ dump here; the downstream - // `dec.audio_injected` (152, hidden) is the real input to the LM and - // is dumped on both sides. + // No clean C++ intermediate matches the reference's un-expanded (15, hidden) + // dec.token_emb; skip it and rely on the downstream dec.audio_injected dump. pb.dumps.token_emb = nullptr; ggml_tensor * x; @@ -260,9 +247,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, return pb; } -// --------------------------------------------------------------------------- // Step graph (single-token decode) -// --------------------------------------------------------------------------- StepBuild build_step_graph(ggml_context * ctx, const CanaryQwenWeights & weights, @@ -356,9 +341,7 @@ StepBuild build_step_graph(ggml_context * ctx, return sb; } -// --------------------------------------------------------------------------- // Batched prefill / step (offline transcribe_run_batch) -// --------------------------------------------------------------------------- PrefillBuildBatched build_prefill_graph_batched( ggml_context * ctx, @@ -392,10 +375,7 @@ PrefillBuildBatched build_prefill_graph_batched( pb.input_ids_in = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, T_prompt_max, B); ggml_set_input(pb.input_ids_in); - // Audio injection is an elementwise blend over the flat token axis (no - // set_rows): a k-quant token_embd get_rows is unsupported on CUDA and runs - // on CPU; a set_rows consuming that CPU tensor straddles the CPU/CUDA split - // and faults. x*keep_mask + audio_dense crosses the split cleanly. + // Audio injection via elementwise blend (no set_rows); see decoder.h. pb.audio_dense_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, static_cast(T_prompt_max) * B); ggml_set_input(pb.audio_dense_in); @@ -418,9 +398,8 @@ PrefillBuildBatched build_prefill_graph_batched( ggml_tensor * ids_flat = ggml_reshape_1d(ctx, pb.input_ids_in, static_cast(T_prompt_max) * B); - // x is 2D [hidden, T_prompt_max*B] (contiguous). Blend the audio embeds in - // elementwise: x = x*keep_mask + audio_dense. keep_mask broadcasts over - // hidden (ne[0]). Then reshape to 3D for the batched block stack. + // x = x*keep_mask + audio_dense (keep_mask broadcasts over hidden), then + // reshape to 3D for the batched block stack. ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed.token_w, ids_flat); x = ggml_add(ctx, ggml_mul(ctx, x, pb.keep_mask_in), pb.audio_dense_in); x = ggml_reshape_3d(ctx, x, hidden, T_prompt_max, B); diff --git a/src/arch/canary_qwen/decoder.h b/src/arch/canary_qwen/decoder.h index 86d23633..cd4082cc 100644 --- a/src/arch/canary_qwen/decoder.h +++ b/src/arch/canary_qwen/decoder.h @@ -83,11 +83,11 @@ StepBuild build_step_graph(ggml_context * ctx, struct PrefillBuildBatched { ggml_tensor * input_ids_in = nullptr; // [T_prompt_max, B] i32 - // Audio injection via elementwise blend (no set_rows): audio_dense holds - // the audio embeds scattered (host-side) into their prompt positions, zero - // elsewhere; keep_mask is 0 at audio positions and 1 elsewhere. The block - // input is x*keep_mask + audio_dense. Elementwise ops cross the CPU/CUDA - // split (forced by k-quant token_embd get_rows) cleanly, unlike a set_rows. + // Audio injection via elementwise blend (no set_rows): audio_dense holds the + // audio embeds scattered host-side into their prompt positions (zero + // elsewhere), keep_mask is 0 there and 1 elsewhere, block input is + // x*keep_mask + audio_dense. Elementwise ops cross the CPU/CUDA split + // (forced by k-quant token_embd get_rows) cleanly, unlike a set_rows. ggml_tensor * audio_dense_in = nullptr; // [hidden, T_prompt_max*B] f32 ggml_tensor * keep_mask_in = nullptr; // [1, T_prompt_max*B] f32 (0=audio,1=keep) ggml_tensor * positions_in = nullptr; // [T_prompt_max] i32 diff --git a/src/arch/canary_qwen/model.cpp b/src/arch/canary_qwen/model.cpp index b6464d97..df856f60 100644 --- a/src/arch/canary_qwen/model.cpp +++ b/src/arch/canary_qwen/model.cpp @@ -14,11 +14,8 @@ // -> Qwen3-1.7B causal LM (28 blocks) // -> tied lm_head -> greedy autoregressive loop until EOS or max_new. // -// BF16 weight promotion: the canary-1b-flash perception encoder was -// validated F32-only; the SALM checkpoint stores its weights as BF16. -// On CPU primary backend we promote ALL BF16 linear weights (encoder -// and decoder) to F32 at load time so ggml's CPU matmul hits the F32 -// path. This matches the reference's `model_dtype = f32` regime. +// On CPU primary backend all BF16 linear weights are promoted to F32 at load +// (matches the reference's f32 regime); see promote_linears_bf16_to_f32_on_cpu. #include "canary_qwen.h" @@ -62,10 +59,6 @@ extern const Arch arch; static_assert(std::is_base_of_v); static_assert(std::is_base_of_v); -// --------------------------------------------------------------------------- -// Lifecycle -// --------------------------------------------------------------------------- - CanaryQwenSession::~CanaryQwenSession() { kv_cache.free(); kv_cache_batch.free(); @@ -128,20 +121,15 @@ namespace { constexpr const char k_default_variant[] = "canary-qwen-2.5b"; constexpr float kBnEps = 1e-5f; -// --------------------------------------------------------------------------- // Input-length contract (see docs/input-limits.md). canary_qwen is a // hard-context-cap family: audio tokens + chat prompt + generation share the -// Qwen3 decoder's context window (decoder.max_position_embeddings). The KV -// cache grows to fit the prompt, clamped to that ceiling (replacing the -// earlier fixed 2048 wall that ignored the model's real context). Over-length -// input is rejected up front with TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript -// that fills the per-run generation budget before end-of-stream is flagged via -// transcribe_was_truncated(). -// --------------------------------------------------------------------------- - -// Per-run generation budget. Matches the `max_new` literal used in both the -// single-utterance step loop and the batched step loop below; keep all three -// in sync. +// Qwen3 decoder's context window (decoder.max_position_embeddings), clamped to +// that ceiling. Over-length input is rejected with +// 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; // Effective decoder context ceiling, in tokens: the model's trained maximum @@ -186,25 +174,13 @@ int64_t canary_qwen_max_audio_ms(const CanaryQwenHParams & hp) { return mel_frames * hp.fe_hop_length * 1000 / hp.fe_sample_rate; } -// --------------------------------------------------------------------------- -// Chat-template prompt construction -// --------------------------------------------------------------------------- -// -// Per the reference SALM trace, HF's chat template renders the JFK -// request as the literal token sequence: -// <|im_start|> user \n Trans cribe the following : -// <|audioplaceholder|> -// <|im_end|> \n <|im_start|> assistant \n -// -// Verified empirically: tok.encode("user\nTranscribe the following: ") -// -> [872, 198, 3167, 3114, 279, 2701, 25, 220]. -// -// We pre-encode the prefix ("<|im_start|>" + "user\nTranscribe the -// following: ") and suffix ("<|im_end|>\n<|im_start|>assistant\n") at -// load time. The audio_locator id is replicated T_enc times between -// them at run time. -// -// The ChatTokens struct holds the few special-token ids we need. +// Chat-template prompt construction. HF's chat template renders the request as: +// <|im_start|> user \n Transcribe the following: <|audioplaceholder|> +// <|im_end|> \n <|im_start|> assistant \n +// Verified: tok.encode("user\nTranscribe the following: ") +// -> [872, 198, 3167, 3114, 279, 2701, 25, 220]. We pre-encode prefix and +// suffix at load; the audio_locator id is replicated T_enc times between them +// at run time. transcribe_status resolve_chat_tokens(const transcribe::Tokenizer & tok, ChatTokens & out) @@ -271,11 +247,8 @@ transcribe_status build_static_prompt_segments(CanaryQwenModel & m) { return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// BatchNorm fusion (same math as canary). Pre-fused (scale, bias) fall -// out of (running_mean, running_var, weight, bias, eps). -// --------------------------------------------------------------------------- - +// BatchNorm fusion (same math as canary): pre-fused (scale, bias) fall out of +// (running_mean, running_var, weight, bias, eps). transcribe_status fuse_batch_norm(CanaryQwenModel & m) { const size_t n_blocks = m.weights.blocks.size(); if (n_blocks == 0) return TRANSCRIBE_OK; @@ -321,31 +294,14 @@ transcribe_status fuse_batch_norm(CanaryQwenModel & m) { return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // F16 → F32 promotion for pointwise AND depthwise conv kernels (any backend). -// -// canary-1b-flash (the closest sibling) ships F32 conv kernels because its -// reference dtype is F32. canary-qwen-2.5b is BF16-reference, so the -// converter stores conv kernels at F16 (the loader rejects BF16 conv). -// -// `ggml_conv_2d_dw_direct` silently produces wildly wrong output for F16 -// kernels (verified on both CPU and Metal): the conv module's output -// magnitude blows up by ~1500x on the first block, destroying every -// downstream value (decoder collapses to a single `!` token). Promoting -// the depthwise kernel to F32 at load time bypasses that path's dtype -// handling and matches NeMo to single-percent drift. -// -// The pointwise kernels were originally promoted only on CPU (a perf fix -// — the F16 path is slow there because of dequant inside matmul) but -// promoting them on Metal too is harmless and simplifies the policy. -// -// We do NOT call `load_common::promote_conv_pw_f16_to_f32_on_cpu` because -// that helper has a hard `if (primary_kind != Cpu) return OK` early-out. -// Refactoring the shared helper to accept a force-flag is the right -// long-term move; for now we duplicate the logic here so the canary_qwen -// fix does not depend on cross-family churn. -// --------------------------------------------------------------------------- - +// The converter stores conv kernels at F16 (canary-qwen-2.5b is BF16-reference; +// the loader rejects BF16 conv). `ggml_conv_2d_dw_direct` silently produces +// wildly wrong output for F16 kernels (verified CPU + Metal): the conv output +// blows up ~1500x on the first block, collapsing the decoder to a single `!`. +// Promoting to F32 bypasses that and matches NeMo to single-percent drift. +// (load_common::promote_conv_pw_f16_to_f32_on_cpu can't be used — it early-outs +// on non-CPU backends.) transcribe_status promote_conv_pw_to_f32_on_cpu(CanaryQwenModel & m) { if (m.plan.primary == nullptr) return TRANSCRIBE_OK; @@ -433,24 +389,10 @@ transcribe_status promote_conv_pw_to_f32_on_cpu(CanaryQwenModel & m) { return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// CPU-only BF16 → F32 promotion of all linear weights. -// -// The reference (NeMo SALM dumper) loads the model in F32 and runs F32 -// inference. The canary-1b-flash encoder (whose weights canary_qwen -// reuses) was validated F32-only when ported; the BF16 matmul path -// through ggml's CPU backend is unverified for this graph topology. -// -// We promote ALL BF16 linear weights to F32 once at load time. This -// stays inside the existing F32 weight buffer (separate from -// backend_buffer, so we can free it independently). The graph builders -// don't need to change — they read whatever dtype each tensor pointer -// reports. -// -// On non-CPU primary backend (Metal, Vulkan), we skip — those backends -// have their own well-tested BF16 paths. -// --------------------------------------------------------------------------- - +// CPU-only BF16 → F32 promotion of all linear weights, matching the reference's +// F32 inference regime. Replacements live in a dedicated buffer (freed +// independently); graph builders read whatever dtype each tensor reports. +// Non-CPU backends (Metal, Vulkan) keep their own BF16 paths. transcribe_status promote_linears_bf16_to_f32_on_cpu(CanaryQwenModel & m) { if (m.plan.primary_kind != transcribe::BackendKind::Cpu || m.plan.primary == nullptr) @@ -565,10 +507,7 @@ transcribe_status promote_linears_bf16_to_f32_on_cpu(CanaryQwenModel & m) { return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // Loader entry points (forward-declared so we can register `arch` after). -// --------------------------------------------------------------------------- - extern transcribe_status load (Loader &, const transcribe_model_load_params *, transcribe_model **); extern transcribe_status init_context(transcribe_model *, const transcribe_session_params *, @@ -576,10 +515,6 @@ extern transcribe_status init_context(transcribe_model *, const transcribe_sessi extern transcribe_status run (transcribe_session *, const float *, int, const transcribe_run_params *); -// --------------------------------------------------------------------------- -// load -// --------------------------------------------------------------------------- - transcribe_status load( Loader & loader, const transcribe_model_load_params * params, @@ -606,12 +541,11 @@ transcribe_status load( st != TRANSCRIBE_OK) return st; // Publish the input-length ceiling now that the decoder context window - // and frontend rate are known. See docs/input-limits.md. + // and frontend rate are known. m->caps.max_audio_ms = canary_qwen_max_audio_ms(m->hparams); - // Basis for the session-level limits query (transcribe_session_get_limits): - // the same constants canary_qwen_max_audio_ms uses, kept so the effective - // limit can be recomputed at a lowered session n_ctx. + // Basis for transcribe_session_get_limits: the same constants + // canary_qwen_max_audio_ms uses, so the limit recomputes at a lowered n_ctx. if (m->hparams.dec_max_position > 0 && m->hparams.enc_subsampling_factor > 0 && m->hparams.fe_hop_length > 0 && m->hparams.fe_sample_rate > 0) { @@ -648,7 +582,7 @@ transcribe_status load( if (auto st = resolve_chat_tokens(m->tok, m->chat_tokens); st != TRANSCRIBE_OK) return st; if (auto st = build_static_prompt_segments(*m); st != TRANSCRIBE_OK) return st; - // ---- Mel frontend ---- + // Mel frontend. { transcribe::MelConfig cfg {}; cfg.sample_rate = m->hparams.fe_sample_rate; @@ -689,7 +623,7 @@ transcribe_status load( m->mel.emplace(cfg); } - // ---- Reopen GGUF with no_alloc to bind the tensor catalog ---- + // Reopen GGUF with no_alloc to bind the tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -704,7 +638,7 @@ transcribe_status load( return st; } - // ---- Backend plan + alloc + stream tensor data ---- + // Backend plan + alloc + stream tensor data. const transcribe_backend_request backend_req = (params != nullptr) ? params->backend : TRANSCRIBE_BACKEND_AUTO; if (auto st = load_common::init_backends(backend_req, (params != nullptr) ? params->gpu_device : 0, "canary_qwen", m->plan); @@ -736,14 +670,9 @@ transcribe_status load( } gguf_free(gguf_data); - // ---- Post-load weight transformations ---- + // Post-load weight transformations. if (auto st = fuse_batch_norm(*m); st != TRANSCRIBE_OK) return st; - - // F16 conv pointwise → F32 (CPU only). Same as canary. if (auto st = promote_conv_pw_to_f32_on_cpu(*m); st != TRANSCRIBE_OK) return st; - - // BF16 linears → F32 (CPU only). Mitigates the canary-1b-flash - // F32-only validation gap for this BF16 GGUF. if (auto st = promote_linears_bf16_to_f32_on_cpu(*m); st != TRANSCRIBE_OK) return st; // Pack gate + up into one tensor per layer (one mul_mat instead of two). @@ -771,10 +700,6 @@ transcribe_status load( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// init_context -// --------------------------------------------------------------------------- - transcribe_status init_context( transcribe_model * model, const transcribe_session_params * params, @@ -788,13 +713,10 @@ transcribe_status init_context( cc->kv_type = params->kv_type; cc->n_ctx = transcribe_session_params_n_ctx(params); - // Reference runs without flash; default to off for tightest tensor - // parity. TRANSCRIBE_FLASH_DECODER=1 (or =encoder) overrides. - // Flash defaults: encoder off (the FastConformer rel-pos path has a - // manual rel_shift trick that ggml_flash_attn_ext doesn't subsume), - // decoder ON (the Qwen3 LM step graph is dispatch-bound on Metal — - // turning flash on takes per-step kernel count from ~22 ops/layer to - // ~14 and yields ~2.4x decode speedup measured on jfk.wav). + // Flash defaults: encoder off (the FastConformer rel-pos path has a manual + // rel_shift trick ggml_flash_attn_ext doesn't subsume), decoder ON (~2.4x + // decode speedup on jfk.wav, the Qwen3 step graph is dispatch-bound on + // Metal). TRANSCRIBE_FLASH_DECODER=1 (or =encoder) overrides. cc->encoder_use_flash = false; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, @@ -825,10 +747,6 @@ transcribe_status init_context( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// run -// --------------------------------------------------------------------------- - void apply_thread_policy(CanaryQwenSession * cc) { transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } @@ -903,7 +821,7 @@ transcribe_status run(transcribe_session * context, } cc->t_mel_us = ggml_time_us() - t_mel_start; - // ---- Encoder ---- + // Encoder. if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; @@ -944,7 +862,6 @@ transcribe_status run(transcribe_session * context, return TRANSCRIBE_ERR_OOM; } - // Upload mel. ggml_backend_tensor_set(eb.mel_in, cc->mel_buf.data(), 0, cc->mel_buf.size() * sizeof(float)); if (transcribe::debug::enabled()) { @@ -1001,7 +918,7 @@ transcribe_status run(transcribe_session * context, try_dump("enc.final", eb.dumps.final_out, "encoder.final"); try_dump("perception.proj.out", eb.dumps.perception_out, "perception.proj.out"); - // ---- Read perception output to host ---- + // Read perception output to host. const int hidden = static_cast(eb.out->ne[0]); const int T_enc = static_cast(eb.out->ne[1]); cc->enc_host.resize(static_cast(hidden) * @@ -1009,7 +926,7 @@ transcribe_status run(transcribe_session * context, ggml_backend_tensor_get(eb.out, cc->enc_host.data(), 0, cc->enc_host.size() * sizeof(float)); - // ---- Build prompt token-id list ---- + // Build prompt token-id list. std::vector prompt_ids; prompt_ids.reserve(cm->prompt_prefix_ids.size() + T_enc + cm->prompt_suffix_ids.size()); @@ -1026,11 +943,8 @@ transcribe_status run(transcribe_session * context, const int T_prompt = static_cast(prompt_ids.size()); const int suffix_len = T_prompt - prefix_len - T_enc; - // ---- Input-length gate (see docs/input-limits.md) ---- - // audio tokens + prompt + generation must fit the decoder context window - // (optionally lowered by the caller's n_ctx). The audio-token count is - // fixed by the input length, so reject an over-length clip here, before - // prefill/decode, instead of walling at a fixed size. + // 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) { transcribe::log_msg( @@ -1043,13 +957,10 @@ transcribe_status run(transcribe_session * context, return TRANSCRIBE_ERR_INPUT_TOO_LONG; } - // ---- KV cache init (grow-to-fit, clamped to the context ceiling) ---- - // Size to hold the prompt plus the generation budget, rounded up to a - // 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. For typical - // short clips this is the same 1024/2048 width as before, so the decode is - // unchanged — only previously-walled long clips now fit. + // 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 + // 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) want_n_ctx *= 2; if (want_n_ctx > ceiling) want_n_ctx = ceiling; @@ -1202,7 +1113,7 @@ transcribe_status run(transcribe_session * context, int32_t next_tok = argmax(logits); generated_ids.push_back(next_tok); - // ---- Step loop ---- + // Step loop. const int32_t eos_id = hp.eos_token_id; const int max_new = k_max_new; int cur_past = T_prompt; @@ -1264,10 +1175,9 @@ transcribe_status run(transcribe_session * context, const ggml_fp16_t mask_neg_inf = ggml_fp32_to_fp16(-INFINITY); std::vector step_mask(max_n_kv, mask_neg_inf); - // Mid-generation tensor coverage: `dec.logits_raw.gen8` is the - // logits for the 9th lm_head call (= step iter 7 of our step loop; - // prefill = 1st call, iter K = (K+2)th call). Same semantics as - // funasr_nano's gen_dump_step. + // Mid-generation tensor coverage: `dec.logits_raw.gen8` is the logits for + // the 9th lm_head call (= step iter 7; prefill = 1st call, iter K = + // (K+2)th call). constexpr int gen_dump_step = 7; int n_steps = 0; @@ -1330,10 +1240,8 @@ transcribe_status run(transcribe_session * context, n_steps += 1; } - // The decode stopped at EOS (complete) or at the generation budget / - // context width (truncated). Surface the latter via - // transcribe_was_truncated() and a WARN rather than returning a silently - // shortened transcript. See docs/input-limits.md. + // Decode stopped at EOS (complete) or the generation budget / context width + // (truncated). Surface the latter via transcribe_was_truncated() + WARN. if (next_tok != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -1399,9 +1307,8 @@ transcribe_status run(transcribe_session * context, / static_cast(hp.fe_sample_rate); cc->segments.push_back(std::move(seg)); - // The partial transcript is fully populated above; a truncated decode - // returns the hard OUTPUT_TRUNCATED status (the result stays readable, - // like an aborted run). See docs/input-limits.md. + // A truncated decode returns OUTPUT_TRUNCATED; the partial transcript above + // stays readable (like an aborted run). return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } @@ -1412,9 +1319,9 @@ transcribe_status run(transcribe_session * context, // Offline batched decode (transcribe_run_batch) // =========================================================================== // Serial mel + conformer encoder (incl. perception projection) per utterance -// produce each one's audio embedding [hidden, T_enc]; then the prefill and the -// autoregressive step loop are batched via the shared causal_lm primitives — -// the same recipe as arch/qwen3_asr and arch/funasr_nano. +// produce each one's audio embedding [hidden, T_enc]; the prefill and +// autoregressive step loop are then batched via the shared causal_lm +// primitives. namespace { @@ -1431,9 +1338,8 @@ transcribe_status reset_ctx(CanaryQwenSession * cc, int mb) { return cc->compute_ctx != nullptr ? TRANSCRIBE_OK : TRANSCRIBE_ERR_GGUF; } -// mel + conformer encoder + perception for one utterance → [hidden, T_enc]. -// Encoder (+ perception) for one utterance from a PRECOMPUTED mel buffer -// (the mel is computed in parallel by the caller) → [hidden, T_enc]. +// Conformer encoder (+ perception) for one utterance from a PRECOMPUTED mel +// buffer (the mel is computed in parallel by the caller) → [hidden, T_enc]. transcribe_status encode_one( CanaryQwenSession * cc, CanaryQwenModel * cm, const std::vector & mel_buf, int mel_n_frames, @@ -1522,17 +1428,13 @@ transcribe_status run_batch( transcribe::debug::init(); const auto & hp = cm->hparams; - // Decoder context ceiling for the input-length gate (see - // docs/input-limits.md). Same value the single-utterance run() enforces; - // an over-length utterance is rejected with TRANSCRIBE_ERR_INPUT_TOO_LONG - // rather than walling at a fixed KV size. + // Decoder context ceiling for the input-length gate. Same value the + // single-utterance run() enforces. const int ceiling = canary_qwen_context_ceiling(cc->n_ctx, hp); - // ---- Pass 1: per-utterance mel + encoder (serial) ---- std::vector valid(n, 0); - // Per-utterance failure status for the result capture below. Defaults to - // INVALID_ARG (bad pcm / mel / encode); the input-length gate upgrades it - // to INPUT_TOO_LONG for over-length clips. + // Per-utterance failure status. Defaults to INVALID_ARG (bad pcm/mel/ + // encode); the input-length gate upgrades it to INPUT_TOO_LONG. std::vector fail_status(n, TRANSCRIBE_ERR_INVALID_ARG); std::vector> enc_hosts(n); std::vector T_enc(n, 0); @@ -1541,7 +1443,7 @@ transcribe_status run_batch( int prefix_len = static_cast(cm->prompt_prefix_ids.size()); int64_t mel_us = 0, enc_us = 0; - // ---- Pass 0: parallel mel (MelFrontend::compute is thread-safe) ---- + // Pass 0: parallel mel (MelFrontend::compute is thread-safe). std::vector> mel_bufs(n); std::vector mel_nf(n, 0); int n_mel_threads = cc->n_threads; @@ -1558,7 +1460,7 @@ transcribe_status run_batch( }); mel_us += ggml_time_us() - t_mel0; - // ---- Pass 1: per-utterance encoder (serial) ---- + // Pass 1: per-utterance encoder (serial). for (int b = 0; b < n; ++b) { if (cc->poll_abort()) return TRANSCRIBE_ERR_ABORTED; if (mel_nf[b] <= 0) continue; @@ -1574,9 +1476,8 @@ transcribe_status run_batch( cm->prompt_suffix_ids.end()); T_prompt[b] = static_cast(prompt_ids[b].size()); - // Input-length gate (see docs/input-limits.md). audio tokens + - // prompt + generation must fit the decoder context window; reject an - // over-length utterance here instead of walling at a fixed KV size. + // 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) { transcribe::log_msg( TRANSCRIBE_LOG_LEVEL_ERROR, @@ -1608,12 +1509,11 @@ transcribe_status run_batch( const int max_new = k_max_new; int max_n_kv = 1024; while (max_n_kv < max_T_prompt + max_new) max_n_kv *= 2; - // Clamp to the context ceiling. The per-utterance gate above guarantees - // max_T_prompt + max_new <= ceiling for every valid utterance, so this - // never truncates an in-spec workload — it only bounds the pow2 round-up. + // Clamp the pow2 round-up to the ceiling (the per-utterance gate guarantees + // every valid row still fits). if (max_n_kv > ceiling) max_n_kv = ceiling; - // ---- Allocate batched KV cache ---- + // Allocate batched KV cache. ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; if (cc->kv_cache_batch.self_k == nullptr || @@ -1634,7 +1534,7 @@ transcribe_status run_batch( ggml_backend_buffer_clear(cc->kv_cache_batch.buffer, 0); } - // ---- Pass 2: batched prefill ---- + // Pass 2: batched prefill. std::vector next_tok(n, 0); std::vector n_past(n, 0); std::vector> generated(n); @@ -1649,9 +1549,9 @@ transcribe_status run_batch( const int hidden = hp.dec_hidden; std::vector ids(static_cast(max_T_prompt) * n, 0); - // Audio injection by elementwise blend: audio_dense holds each - // utterance's audio embeds scattered into their prompt positions (zero - // elsewhere), keep is 0 there and 1 elsewhere. x = x*keep + audio_dense. + // Audio injection by elementwise blend (see decoder.h): audio_dense + // holds each utterance's audio embeds scattered into their prompt + // positions, keep is 0 there and 1 elsewhere. std::vector audio_dense(static_cast(hidden) * max_T_prompt * n, 0.0f); std::vector keep(static_cast(max_T_prompt) * n, 1.0f); std::vector kidx(static_cast(max_T_prompt) * n); @@ -1709,7 +1609,7 @@ transcribe_status run_batch( } } - // ---- Pass 3: batched step loop (shared causal_lm driver) ---- + // Pass 3: batched step loop (shared causal_lm driver). const int32_t eos_id = cm->hparams.eos_token_id; if (reset_ctx(cc, 16) != TRANSCRIBE_OK) return TRANSCRIBE_ERR_GGUF; @@ -1742,7 +1642,7 @@ transcribe_status run_batch( } const int64_t step_us = step_stats.step_us; - // ---- Capture results ---- + // Capture results. const int valid_count = std::max(1, static_cast( std::count(valid.begin(), valid.end(), char(1)))); for (int b = 0; b < n; ++b) { @@ -1763,10 +1663,8 @@ transcribe_status run_batch( seg.t1_ms = static_cast(n_samples[b]) * 1000 / static_cast(hp.fe_sample_rate); rs.segments.push_back(std::move(seg)); - // Per-utterance truncation parity with single-shot run(): a row cut at - // the generation budget / KV window before eos reports - // TRANSCRIBE_ERR_OUTPUT_TRUNCATED (partial transcript retained). Only - // override a TRANSCRIBE_OK status, never a worse one. + // Per-utterance truncation parity with single-shot run(). Only override + // a TRANSCRIBE_OK status, never a worse one. if (b < static_cast(truncated.size()) && truncated[b] && rs.status == TRANSCRIBE_OK) { cc->was_truncated = true; diff --git a/src/arch/cohere/capabilities.cpp b/src/arch/cohere/capabilities.cpp index e5eaf8b9..8a8bb88e 100644 --- a/src/arch/cohere/capabilities.cpp +++ b/src/arch/cohere/capabilities.cpp @@ -10,19 +10,10 @@ void apply_family_invariants(transcribe_model & model) { caps.native_sample_rate = 16000; caps.supports_translate = false; - // Cohere ASR uses an autoregressive decoder with the - // <|notimestamp|> prompt token hard-wired on every run, so the - // generated token stream carries no alignment information. The - // family emits a full-text transcript inside a single segment - // with zeroed t0/t1 — there is no honest token, word, or - // segment timing to return. Advertise NONE and let the - // dispatcher reject any request for finer granularity. + // <|notimestamp|> is hard-wired, so the token stream carries no + // alignment info: emit text-only with zeroed timings, advertise NONE. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; - // Cancellation is wired at the per-run level. The other unallied - // features (initial prompt, temperature fallback, long-form, PNC, - // ITN) are Whisper-specific or family-specific elsewhere; Cohere - // ASR does not participate. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); } diff --git a/src/arch/cohere/cohere.h b/src/arch/cohere/cohere.h index caaa5f94..79b64f55 100644 --- a/src/arch/cohere/cohere.h +++ b/src/arch/cohere/cohere.h @@ -34,16 +34,9 @@ namespace transcribe::cohere { void apply_family_invariants(transcribe_model & model); -// --------------------------------------------------------------------------- -// KV cache for the autoregressive decoder. -// -// Follows the whisper.cpp pattern: flat 1D tensors for K and V, with -// views used during graph construction to read/write per-layer slices. -// -// Self-attention cache: grows with each decode step. -// Cross-attention cache: computed once from encoder output, then reused. -// --------------------------------------------------------------------------- - +// KV cache for the autoregressive decoder. Flat 1D K/V tensors (whisper.cpp +// pattern), per-layer slices via views. Self cache grows per step; +// cross cache is computed once from encoder output, then reused. struct CohereKvCache { // Self-attention KV cache. // Flat tensors of size [n_state * n_layer * n_ctx]. @@ -137,11 +130,8 @@ struct CohereModel final : public transcribe_model { CohereWeights weights; ggml_context * ctx_meta = nullptr; - // Runtime backend plan — see transcribe-backend.h. Replaces the - // old `std::vector` field: the plan's scheduler_list - // holds the same handles, plus a classified primary kind so - // helpers don't have to re-derive it via ggml_backend_name - // string matching. + // Runtime backend plan — see transcribe-backend.h. scheduler_list + // holds the backend handles plus a classified primary kind. transcribe::BackendPlan plan; ggml_backend_buffer_t backend_buffer = nullptr; @@ -150,10 +140,8 @@ struct CohereModel final : public transcribe_model { ggml_backend_buffer_t bn_fused_buffer = nullptr; // On CPU primary backend, the conformer 1×1 pointwise conv weights - // are dequantized from their on-disk F16 form to F32 at load time - // — Zen 2 class CPUs pay an F16→F32 upconvert on every matmul that - // erases the bandwidth win. Tensors live in this ctx, and the - // CohereBlock slots point here instead of the main weight buffer. + // 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; @@ -179,21 +167,10 @@ struct CohereSession final : public transcribe_session { std::vector enc_host; - // Flash-attention is controlled per-stage because the encoder and - // decoder have different head dimensions and therefore different - // backend support profiles: - // - // - encoder head_dim = 160 -> upstream ggml's Metal backend has - // no flash_attn_ext kernel for this dk, so we default it OFF - // on Metal. Manual mul_mat + softmax + mul_mat ties or beats - // flash at encoder sequence lengths anyway. - // - decoder head_dim = 128 -> works with flash_attn_ext on - // every backend we ship today, so we default it ON. - // - // The TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH env vars apply - // to both stages at once (the user's intent is "no flash kernels - // anywhere" or "flash kernels everywhere"). Backend-specific auto - // disable is per-stage. + // Flash-attn is per-stage: encoder dk=160 has no Metal flash_attn_ext + // kernel (default OFF on Metal; manual path ties anyway); decoder + // dk=128 works everywhere (default ON). TRANSCRIBE_NO_FLASH / + // TRANSCRIBE_FORCE_FLASH apply to both stages at once. bool encoder_use_flash = true; bool decoder_use_flash = true; diff --git a/src/arch/cohere/decoder.cpp b/src/arch/cohere/decoder.cpp index 01eec367..77a437f7 100644 --- a/src/arch/cohere/decoder.cpp +++ b/src/arch/cohere/decoder.cpp @@ -1,29 +1,15 @@ // arch/cohere/decoder.cpp - Cohere ASR Transformer decoder graph builder. // -// Implements three graph builders: +// Graph builders: build_decoder_graph (non-cached prompt pass, kept for +// reference), build_cross_kv_graph (cross-attn K/V once per utterance), +// build_decoder_graph_kv (KV-cached, both prompt and step passes), plus the +// static-topology step graphs. // -// 1. build_decoder_graph() -- Original prompt-pass graph (no KV cache). -// Kept for reference/validation. Processes all tokens at once with -// a causal self-attention mask. -// -// 2. build_cross_kv_graph() -- Computes cross-attention K/V for all -// decoder layers from the encoder output. Run once per utterance. -// Writes into the cross-attention KV cache via ggml_cpy. -// -// 3. build_decoder_graph_kv() -- KV-cached decoder graph. Works for -// both prompt pass (n_tokens > 1) and step pass (n_tokens = 1). -// Reads/writes the self-attention KV cache; reads cross-attention -// KV cache. -// -// The decoder forward pass: -// 1. Token embedding lookup + positional embedding lookup + LayerNorm -// 2. For each layer: -// a. Pre-LN -> causal self-attention -> residual add -// b. Pre-LN -> cross-attention (encoder as K/V) -> residual add -// c. Pre-LN -> FFN (ReLU) -> residual add -// 3. Final LayerNorm -// 4. Linear projection (tied weight = token embedding^T) + bias -// 5. Log-softmax +// Decoder forward pass: +// 1. token + positional embedding lookups + LayerNorm +// 2. per layer: Pre-LN self-attn / cross-attn / FFN, each with residual add +// 3. final LayerNorm +// 4. tied LM head (token embedding^T) + bias + optional log-softmax #include "decoder.h" @@ -197,10 +183,8 @@ ggml_tensor * mha_self_cached( Q = ggml_permute(ctx, Q, 0, 2, 1, 3); // Q is now [head_dim, n_tokens, n_heads, 1] - // Write new K/V into cache. - // Cache layout (flash_attn mode): K and V are stored contiguously - // per-layer as [hidden, n_ctx] (i.e., n_state elements per position). - // Layer il occupies offset il*n_ctx*hidden in the flat 1D buffer. + // Write new K/V into cache. Layout: per-layer [hidden, n_ctx] contiguous + // (n_state per position); layer il at flat offset il*n_ctx*hidden. { ggml_tensor * k_dst = ggml_view_1d(ctx, kv_cache.self_k, static_cast(n_tokens) * hidden, @@ -234,26 +218,20 @@ ggml_tensor * mha_self_cached( ggml_element_size(kv_cache.self_v) * static_cast( static_cast(il) * n_ctx * hidden)); - // Attention. Flash path is fused; manual path is the usual - // mul_mat + soft_max_ext + mul_mat triple. See `mha` above for - // the shape argument. + // Flash path is fused; manual path is the mul_mat + soft_max_ext + + // mul_mat triple (see `mha` above for the shape walkthrough). ggml_tensor * o; if (use_flash) { o = ggml_flash_attn_ext(ctx, Q, K, V, mask, scale, 0.0f, 0.0f); o = ggml_reshape_2d(ctx, o, hidden, n_tokens); } else { - // K is [head_dim, n_kv, n_heads, 1] from the cache view. - // mul_mat(K, Q) -> [n_kv, n_tokens, n_heads, 1] + // Manual path — see `mha` for the shape walkthrough. + // K: [head_dim, n_kv, n_heads, 1]; V transposed for the value matmul. ggml_tensor * kq = ggml_mul_mat(ctx, K, Q); ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, mask, scale, 0.0f); - - // V is [head_dim, n_kv, n_heads, 1]; we need - // [n_kv, head_dim, n_heads, 1] for the value matmul. ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, V, 1, 0, 2, 3)); - - // attn_weights @ V -> [head_dim, n_tokens, n_heads, 1] o = ggml_mul_mat(ctx, v_t, kq_soft); o = ggml_permute(ctx, o, 0, 2, 1, 3); o = ggml_cont(ctx, o); @@ -415,26 +393,20 @@ ggml_tensor * mha_cross_cached( ggml_element_size(kv_cache.cross_v) * static_cast( static_cast(il) * T_enc * hidden)); - // No mask for cross-attention. Flash path is the fused kernel; - // manual path is the usual triple. + // No mask for cross-attention. ggml_tensor * o; if (use_flash) { o = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); o = ggml_reshape_2d(ctx, o, hidden, n_tokens); } else { - // K is [head_dim, T_enc, n_heads, 1] from the cache view. - // mul_mat(K, Q) -> [T_enc, n_tokens, n_heads, 1] + // Manual path — see `mha` for the shape walkthrough. + // K: [head_dim, T_enc, n_heads, 1]; V transposed for the value matmul. ggml_tensor * kq = ggml_mul_mat(ctx, K, Q); ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, nullptr, scale, 0.0f); - - // V is [head_dim, T_enc, n_heads, 1]; transpose for the - // value matmul. ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, V, 1, 0, 2, 3)); - - // attn_weights @ V -> [head_dim, n_tokens, n_heads, 1] o = ggml_mul_mat(ctx, v_t, kq_soft); o = ggml_permute(ctx, o, 0, 2, 1, 3); o = ggml_cont(ctx, o); @@ -449,10 +421,7 @@ ggml_tensor * mha_cross_cached( } // namespace -// --------------------------------------------------------------------------- // Original non-cached prompt-pass graph (kept for reference/validation). -// --------------------------------------------------------------------------- - DecoderBuild build_decoder_graph(ggml_context * ctx, const CohereWeights & w, const CohereHParams & hp, @@ -486,16 +455,14 @@ DecoderBuild build_decoder_graph(ggml_context * ctx, ggml_set_name(db.encoder_out_in, "dec.encoder_out"); ggml_set_input(db.encoder_out_in); - // Token embedding: lookup from [hidden, vocab_size] table. - // ggml_get_rows takes a table of ne=[hidden, vocab_size] and indices - // of ne=[seq_len], producing [hidden, seq_len]. + // Token embedding: get_rows over [hidden, vocab_size] -> [hidden, seq_len]. ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_embed.token_w, db.token_ids_in); tok_emb = named(tok_emb, "dec.token_emb"); ggml_set_output(tok_emb); db.dumps.token_emb = tok_emb; - // Positional embedding: lookup from [hidden, max_seq] table. + // Positional embedding: get_rows over [hidden, max_seq]. ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_embed.pos_enc, db.pos_ids_in); pos_emb = named(pos_emb, "dec.pos_emb"); @@ -509,19 +476,14 @@ DecoderBuild build_decoder_graph(ggml_context * ctx, ggml_set_output(x); db.dumps.embed_norm = x; - // Build causal self-attention mask. Shape [seq_len, seq_len] f16. - // mask[i][j] = 0 if j <= i, else -inf. - // For flash_attn, mask shape should be [T_k, T_q, 1, 1]. + // Causal self-attention mask [T_k=seq_len, T_q=seq_len]: 0 if j <= i else + // -inf. Driver fills the F32 input; cast to F16 for flash_attn. ggml_tensor * causal_mask = nullptr; if (seq_len > 1) { - // ggml_new_tensor_2d(ctx, F32, T_k=seq_len, T_q=seq_len) - // filled with causal pattern, then cast to F16. causal_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, seq_len, seq_len); ggml_set_name(causal_mask, "dec.causal_mask"); ggml_set_input(causal_mask); - // Will be filled by the driver with proper causal mask values. - // Cast to f16 for flash_attn. causal_mask = ggml_cast(ctx, causal_mask, GGML_TYPE_F16); } @@ -579,10 +541,8 @@ DecoderBuild build_decoder_graph(ggml_context * ctx, x = named(x, "dec.out_before_head"); db.dumps.out_before_head = x; - // Head: tied weight (transpose of token embedding) + bias + log-softmax. - // ggml_mul_mat(W, x) where W=[hidden, vocab_size] and x=[hidden, seq_len] - // gives [vocab_size, seq_len]. The token embedding is [hidden, vocab_size] - // which is exactly what we need. + // Head: tied weight (the [hidden, vocab_size] token embedding) + bias + + // log-softmax. mul_mat(token_w, x[hidden,seq]) -> [vocab_size, seq_len]. ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_embed.token_w, x); if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); @@ -609,10 +569,7 @@ DecoderBuild build_decoder_graph(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- // Cross-attention KV computation graph. -// --------------------------------------------------------------------------- - DecoderBuild build_cross_kv_graph(ggml_context * ctx, const CohereWeights & w, const CohereHParams & hp, @@ -647,14 +604,12 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, for (int il = 0; il < hp.dec_n_layers; ++il) { const auto & blk = w.dec_blocks[il]; - // K = cross_k_w * encoder_out + cross_k_b ggml_tensor * Kcross = ggml_mul_mat(ctx, blk.cross_k_w, db.encoder_out_in); if (blk.cross_k_b != nullptr) { Kcross = ggml_add(ctx, Kcross, blk.cross_k_b); } - // V = cross_v_w * encoder_out + cross_v_b ggml_tensor * Vcross = ggml_mul_mat(ctx, blk.cross_v_w, db.encoder_out_in); if (blk.cross_v_b != nullptr) { @@ -679,10 +634,7 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- // KV-cached decoder graph (prompt + step passes). -// --------------------------------------------------------------------------- - DecoderBuild build_decoder_graph_kv(ggml_context * ctx, const CohereWeights & w, const CohereHParams & hp, @@ -716,14 +668,12 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, ggml_set_name(db.pos_ids_in, "dec.pos_ids"); ggml_set_input(db.pos_ids_in); - // Token embedding. ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_embed.token_w, db.token_ids_in); tok_emb = named(tok_emb, "dec.token_emb"); ggml_set_output(tok_emb); db.dumps.token_emb = tok_emb; - // Positional embedding. ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_embed.pos_enc, db.pos_ids_in); pos_emb = named(pos_emb, "dec.pos_emb"); @@ -737,28 +687,18 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, ggml_set_output(x); db.dumps.embed_norm = x; - // Causal mask for self-attention. - // Shape: [n_kv, n_tokens] -- for each new query position, which - // cache positions are visible. - // - // For the prompt pass (n_past=0, n_tokens>1): - // mask[k, q] = 0 if k <= q, else -inf - // - // For step pass (n_tokens=1): - // All n_kv positions are visible (mask is all zeros), so we can - // skip the mask entirely. But for correctness with n_tokens>1 in - // prompt pass, we always create the mask. + // Causal mask [n_kv, n_tokens]: mask[k, q] = 0 if k <= q else -inf. Only + // the prompt pass (n_tokens > 1) needs it; a single step sees all n_kv. ggml_tensor * causal_mask = nullptr; if (n_tokens > 1) { - // Need a mask of shape [n_kv, n_tokens]. causal_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, n_tokens); ggml_set_name(causal_mask, "dec.causal_mask"); ggml_set_input(causal_mask); causal_mask = ggml_cast(ctx, causal_mask, GGML_TYPE_F16); } - // Pre-create the graph (we need it for ggml_build_forward_expand in - // the self-attention cache writes). + // Pre-create the graph: the self-attention cache writes call + // ggml_build_forward_expand on it. db.graph = ggml_new_graph_custom(ctx, 8192, false); if (db.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -851,11 +791,9 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.out = logits; ggml_set_output(db.out); - // When skipping log_softmax (prompt and step passes), add argmax on - // GPU so we only need to read back a single int32 instead of the - // full [vocab_size x n_tokens] logits tensor. For n_tokens > 1 we - // only need the argmax of the last position (the next token to - // generate), so take a view of the last row before argmax. + // When skipping log_softmax, argmax on GPU so we read back a single int32 + // instead of the full logits tensor. For n_tokens > 1 only the last + // position matters, so view its row before argmax. if (skip_log_softmax) { ggml_tensor * last_logits = logits; if (n_tokens > 1) { @@ -879,10 +817,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- // Static-topology single-token step graph (GPU dispatch path). -// --------------------------------------------------------------------------- - StepBuild build_step_graph(ggml_context * ctx, const CohereWeights & w, const CohereHParams & hp, @@ -1001,10 +936,7 @@ StepBuild build_step_graph(ggml_context * ctx, return sb; } -// =========================================================================== -// Offline batched decode (B utterances) -// =========================================================================== - +// Offline batched decode (B utterances). namespace { // Batched self-attention step. x: [hidden, B] (one token per utterance). diff --git a/src/arch/cohere/decoder.h b/src/arch/cohere/decoder.h index 760bd139..2908d038 100644 --- a/src/arch/cohere/decoder.h +++ b/src/arch/cohere/decoder.h @@ -52,15 +52,11 @@ struct DecoderBuild { ggml_cgraph * graph = nullptr; }; -// Build a decoder prompt-pass graph. No KV cache -- processes the full +// Build a decoder prompt-pass graph (no KV cache): processes the full // prompt sequence at once with a causal self-attention mask. -// -// seq_len: number of prompt tokens. -// T_enc: number of encoder frames (after enc-dec projection). -// use_flash: true -> fused ggml_flash_attn_ext path; -// false -> manual mul_mat + soft_max_ext + mul_mat path. -// Defaults true because dk=128 has flash kernels on every -// backend we ship. +// seq_len: number of prompt tokens. +// T_enc: number of encoder frames (after enc-dec projection). +// use_flash: fused ggml_flash_attn_ext vs manual mul_mat path. DecoderBuild build_decoder_graph(ggml_context * compute_ctx, const CohereWeights & weights, const CohereHParams & hp, @@ -68,37 +64,21 @@ DecoderBuild build_decoder_graph(ggml_context * compute_ctx, int T_enc, bool use_flash = true); -// Build a graph that computes cross-attention K/V for all decoder -// layers from the encoder output, writing them into the cross-attn -// KV cache. This is run once per utterance. -// -// The graph reads from encoder_out (which must be set as an input) -// and writes into kv_cache.cross_k / kv_cache.cross_v via ggml_cpy. -// -// Returns: a DecoderBuild with only .encoder_out_in and .graph set. +// Compute cross-attention K/V for all decoder layers from encoder_out +// (an input) into kv_cache.cross_k / cross_v via ggml_cpy. Run once per +// utterance. Returns a DecoderBuild with only .encoder_out_in and .graph set. DecoderBuild build_cross_kv_graph(ggml_context * compute_ctx, const CohereWeights & weights, const CohereHParams & hp, CohereKvCache & kv_cache, int T_enc); -// Build a decoder graph that uses the KV cache. Works for both -// prompt pass (n_tokens > 1) and step pass (n_tokens = 1). -// -// n_tokens: number of new tokens to process. -// n_past: number of tokens already in the self-attention KV cache. -// For prompt pass this is 0; for step passes it's the -// running total of previously processed tokens. -// T_enc: number of encoder frames (for cross-attention cache shape). -// -// The graph: -// - Computes Q/K/V for self-attention from the new tokens only. -// - Writes new K/V into the self-attention cache at position n_past -// via ggml_cpy. -// - Reads the full self-attention cache [0..n_past+n_tokens) for -// attention computation. -// - Reads cross-attention K/V from the pre-populated cache. -// - Outputs logits for the new tokens only. +// KV-cached decoder graph for both prompt pass (n_tokens > 1) and step +// pass (n_tokens = 1). Writes new self-K/V at position n_past via ggml_cpy, +// reads the full cache [0..n_past+n_tokens) plus the pre-populated cross +// K/V, outputs logits for the new tokens only. +// n_past: tokens already in the self-attention KV cache (0 for prompt). +// T_enc: number of encoder frames (cross-attention cache shape). DecoderBuild build_decoder_graph_kv(ggml_context * compute_ctx, const CohereWeights & weights, const CohereHParams & hp, @@ -109,17 +89,11 @@ DecoderBuild build_decoder_graph_kv(ggml_context * compute_ctx, bool skip_log_softmax = false, bool use_flash = true); -// Static-topology single-token decoder graph. Built once per utterance -// after the prompt pass and reused for every step in the autoregressive -// loop. Compared to build_decoder_graph_kv with n_tokens=1, this graph -// is independent of n_past: -// - KV cache writes go through ggml_set_rows with kv_idx_in as the -// runtime row index (instead of ggml_cpy at a build-time-baked -// view offset). -// - Self-attn K/V reads span the full [0, max_n_kv) window; mask_in -// gates which positions are valid each step. -// Caller pre-allocates the graph + sched once and only updates the -// runtime inputs (token_id_in, pos_id_in, kv_idx_in, mask_in) per step. +// Static-topology single-token decoder graph, built once per utterance and +// reused for every step. Independent of n_past: KV writes go via +// ggml_set_rows at runtime row index kv_idx_in; self-K/V reads span the full +// [0, max_n_kv) window with mask_in gating valid positions. Caller only +// updates the runtime inputs (token_id_in, pos_id_in, kv_idx_in, mask_in). struct StepBuild { ggml_tensor * token_id_in = nullptr; // i32 [1] ggml_tensor * pos_id_in = nullptr; // i32 [1] @@ -143,9 +117,9 @@ StepBuild build_step_graph(ggml_context * compute_ctx, // --------------------------------------------------------------------------- // Batched cross-attention K/V. encoder_out_in is [hidden, T_enc_max, B] -// (each utterance right-padded to T_enc_max); per-layer K/V are written -// into the batched cross cache slab (layer, b). Padded encoder frames -// produce bias-only K/V that the cross-pad mask discards at attention time. +// (each utterance right-padded to T_enc_max); per-layer K/V written into +// the batched cross cache slab (layer, b). Padded frames produce bias-only +// K/V that the cross-pad mask discards at attention time. DecoderBuild build_cross_kv_graph_batched(ggml_context * compute_ctx, const CohereWeights & weights, const CohereHParams & hp, @@ -153,11 +127,10 @@ DecoderBuild build_cross_kv_graph_batched(ggml_context * compute_ctx, int T_enc_max, int n_batch); -// Static-topology batched single-step graph for B utterances. Reused for -// BOTH the (short, uniform) prompt feed and autoregressive generation: -// each invocation consumes one token per utterance, writes self-KV at -// kv_idx[b], and reads self-KV [0,max_n_kv) (self_mask gates valid slots) -// + cross-KV [0,T_enc_max) (cross_mask gates each utterance's real frames). +// Static-topology batched single-step graph for B utterances, reused for +// both the prompt feed and generation: each call consumes one token per +// utterance, writes self-KV at kv_idx[b], reads self-KV [0,max_n_kv) +// (self_mask) + cross-KV [0,T_enc_max) (cross_mask). struct StepBuildBatched { ggml_tensor * token_ids_in = nullptr; // i32 [B] ggml_tensor * pos_ids_in = nullptr; // i32 [B] diff --git a/src/arch/cohere/encoder.cpp b/src/arch/cohere/encoder.cpp index 5e03a66d..ecec71ff 100644 --- a/src/arch/cohere/encoder.cpp +++ b/src/arch/cohere/encoder.cpp @@ -1,18 +1,9 @@ // arch/cohere/encoder.cpp - Cohere ASR Conformer encoder graph builder. // -// Thin glue over the shared Conformer helpers in src/conformer/. The -// per-op helpers (layer_norm, f32-friendly conv helpers, rel_shift, -// conv_module, rel_pos_mhsa, build_conformer_block, build_pre_encode) -// now live in transcribe::conformer. This file owns only Cohere- -// specific glue: -// -// - detect_direct_dw (different policy from Parakeet: Cohere uses -// the im2col path on Metal and CPU, direct only on Vulkan/CUDA) -// - to_view() helpers that project CoherePreEncode / CohereBlock -// onto the shared nullable-pointer views (FFN, attention, and -// conv biases are all present in Cohere) -// - build_encoder_graph orchestration, including the post-encoder -// enc_dec_proj projection and debug dump-preservation calls +// Thin glue over the shared Conformer helpers in transcribe::conformer. +// Cohere-specific bits: detect_direct_dw (depthwise dispatch policy), +// to_view() projections of CoherePreEncode / CohereBlock onto the shared +// views, and build_encoder_graph (including the enc_dec_proj projection). #include "encoder.h" @@ -34,15 +25,9 @@ namespace { namespace conf = transcribe::conformer; -// ----- Per-family depthwise dispatch policy ----------------------- -// -// Cohere uses the direct ggml_conv_2d_dw_direct path on Vulkan and -// CUDA (both have native kernels), and the im2col + mul_mat path on -// Metal and CPU. This is the opposite of Parakeet on Metal/CPU — the -// divergence is preserved as-is during the extraction to avoid any -// quiet behavior changes. See the ConvPolicy comment in conformer.h -// for why the two sites (in-block dw and pre_encode dw) share the -// same detect result in Cohere. +// Depthwise dispatch policy: direct ggml_conv_2d_dw_direct on Vulkan/CUDA +// (native kernels), im2col + mul_mat on Metal/CPU. See the ConvPolicy +// comment in conformer.h for why the two dw sites share one detect result. bool detect_direct_dw(const char * backend) { const char * env = std::getenv("TRANSCRIBE_CONV_DIRECT_DW"); if (env != nullptr) return true; @@ -192,8 +177,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, for (size_t i = 1; i < w.blocks.size(); ++i) { x = conf::build_conformer_block(ctx, x, eb.pos_emb_in, to_view(w.blocks[i]), bparams); - // Spot-check at the middle block (n_layers/2 - 1 to match - // the reference dump naming). + // Spot-check at the middle block (n_layers/2 - 1). if (i == static_cast(hp.enc_n_layers / 2 - 1)) { char bname[64]; std::snprintf(bname, sizeof(bname), "enc.block.%zu.out", i); @@ -216,8 +200,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, x = conf::named(x, "enc.final"); transcribe::debug::mark_tensor_for_dump(x); - // Encoder-decoder projection. Cohere-specific — not part of the - // shared Conformer block, so it stays here. + // Encoder-decoder projection (Cohere-specific, not in the shared block). { ggml_tensor * proj = ggml_mul_mat(ctx, w.enc_dec_proj.weight, x); proj = ggml_add(ctx, proj, w.enc_dec_proj.bias); diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index 40e54d1d..039fadf2 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -61,10 +61,6 @@ CohereSession::~CohereSession() { encoder_out = nullptr; } -// --------------------------------------------------------------------------- -// KV cache initialization. -// --------------------------------------------------------------------------- - bool kv_cache_init(CohereKvCache & cache, ggml_backend_t backend, int n_ctx, @@ -115,7 +111,6 @@ bool kv_cache_init(CohereKvCache & cache, return false; } - // Zero out the buffers. ggml_backend_buffer_clear(cache.buffer, 0); cache.n_ctx = n_ctx; @@ -233,48 +228,32 @@ namespace { constexpr float kBnEps = 1e-5f; -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Cohere ASR is a -// hard-context-cap family with TWO distinct limits: -// -// (a) INPUT limit — the encoder's relative positional encoding table -// (enc_pos_emb_max_len, 5000 encoder frames for the shipped -// checkpoint). The conformer encoder is *not* unbounded like -// parakeet: encode_positions() is applied at the subsampled -// sequence length T_enc, so T_enc must stay within the trained -// pos-emb span or the runtime-generated table aliases past the -// trained range. This is the binding INPUT bound; it is gated up -// front, before the encoder graph is built. -// -// (b) DECODER self-KV — dec_max_seq (1024) and a 512 max-new-tokens -// cap. These bound the *output*: a transcript that runs long -// enough to exhaust the budget is kept as a partial result and -// flagged via transcribe_was_truncated(), not rejected. -// -// The encoder is the input limit, so transcribe_capabilities::max_audio_ms -// is derived from (a). -// --------------------------------------------------------------------------- +// Input-length contract (canonical reference; see docs/input-limits.md). +// Cohere ASR has TWO distinct limits: +// (a) INPUT — the encoder's relative positional-encoding table +// (enc_pos_emb_max_len). T_enc must stay within the trained pos-emb +// span or the runtime table aliases past the trained range. Binding +// input bound, gated up front; drives caps.max_audio_ms. +// (b) DECODER self-KV — dec_max_seq and a 512 max-new-tokens cap. Bound +// the *output*: an over-budget transcript is kept as a partial result +// and flagged via transcribe_was_truncated(), not rejected. // Predicted encoder frame count T_enc for a given mel frame count. The -// FastConformer pre-encode (build_pre_encode in src/conformer/conformer.cpp) -// downsamples time by enc_subsampling_factor via three stride-2, kernel-3, -// pad-(k-1)/2 convs; each one maps T_in -> floor((T_in - 1) / 2) + 1. We -// fold that exact per-conv recurrence here so the prediction matches the -// graph's T_enc (conformer.cpp documents the net result as floor(T_mel/8)). -// This is a pure host-side count; it touches no graph math. +// FastConformer pre-encode downsamples time via three stride-2, kernel-3, +// pad-(k-1)/2 convs, each mapping T_in -> floor((T_in - 1)/2) + 1. We fold +// that exact recurrence here so the prediction matches the graph's T_enc +// (net result floor(T_mel/8)). Pure host-side count; no graph math. int cohere_predict_t_enc(int mel_n_frames, int subsampling_factor) { if (mel_n_frames <= 0 || subsampling_factor <= 0) { return 0; } - // subsampling_factor 8 == three stride-2 stages. Derive the stage - // count from the factor (log2) so a future geometry change stays - // honest; fall back to the documented 8x if it is not a power of two. + // Derive the stride-2 stage count from the factor (log2); fall back to + // floor(T_mel / factor) if it is not a power of two. int stages = 0; for (int f = subsampling_factor; f > 1; f >>= 1) { ++stages; } if ((1 << stages) != subsampling_factor) { - // Non-pow2 factor: use the documented floor(T_mel / factor). return mel_n_frames / subsampling_factor; } int t = mel_n_frames; @@ -287,17 +266,11 @@ int cohere_predict_t_enc(int mel_n_frames, int subsampling_factor) { return t; } -// Advisory transcribe_capabilities::max_audio_ms: the longest audio whose -// encoder frame count T_enc still fits the trained positional-encoding span -// enc_pos_emb_max_len. Inverts the rate: -// T_enc = mel_frames / subsampling_factor -// mel_frames = ms * sample_rate / (hop_length * 1000) -// => ms = T_enc * subsampling_factor * hop_length * 1000 / sr -// at T_enc == enc_pos_emb_max_len. Returns 0 ("unknown / unbounded") if any -// rate hparam is missing, so a misconfigured model is never advertised with -// a wrong finite number. The decoder self-KV / max-new cap bounds the -// OUTPUT, not the input, so it does not enter this figure (a long-but-fitting -// clip may still truncate its transcript — see transcribe_was_truncated). +// transcribe_capabilities::max_audio_ms: longest audio whose T_enc still +// fits the trained pos-emb span enc_pos_emb_max_len. Inverts the rate: +// ms = T_enc * subsampling_factor * hop_length * 1000 / sr +// at T_enc == enc_pos_emb_max_len. Returns 0 (unknown) if any rate hparam +// is missing, so a misconfigured model is never advertised with a wrong value. int64_t cohere_max_audio_ms(const CohereHParams & hp) { if (hp.enc_pos_emb_max_len <= 0 || hp.enc_subsampling_factor <= 0 || hp.fe_hop_length <= 0 || hp.fe_sample_rate <= 0) { @@ -364,14 +337,12 @@ transcribe_status fuse_batch_norm(CohereModel & m) { return TRANSCRIBE_OK; } -// Fold each encoder layer's Q bias into its pos_bias_u / pos_bias_v -// tensors at load time. In rel_pos_mhsa the math is +// Fold each encoder layer's Q bias into its pos_bias_u / pos_bias_v at +// load time. In rel_pos_mhsa: // q_u = (W_q x + q_b) + pos_u = W_q x + (q_b + pos_u) // q_v = (W_q x + q_b) + pos_v = W_q x + (q_b + pos_v) -// so pre-adding q_b into pos_u/pos_v is mathematically identical and -// drops the explicit `q = q + q_b` graph op (48 fewer ops per encoder). -// After fusion we null out attn_q_b so the graph builder naturally -// skips the add via its existing `if (attn_q_b != nullptr)` guard. +// so pre-adding q_b into pos_u/pos_v is identical and drops the explicit +// `q = q + q_b` graph op. We null attn_q_b after so the builder skips the add. transcribe_status fuse_encoder_q_bias(CohereModel & m) { const size_t n_blocks = m.weights.blocks.size(); if (n_blocks == 0) return TRANSCRIBE_OK; @@ -405,10 +376,8 @@ transcribe_status fuse_encoder_q_bias(CohereModel & m) { ggml_backend_tensor_get(b.attn_pos_u, pos_u.data(), 0, nbytes); ggml_backend_tensor_get(b.attn_pos_v, pos_v.data(), 0, nbytes); - // pos_u/v are [head_dim, n_heads] with ne[0]=head_dim. - // q_b is [d_model] = [head_dim*n_heads]. In memory the two - // share the same element layout (same linear offsets), so - // a flat element-wise add is correct. + // pos_u/v [head_dim, n_heads] and q_b [d_model=head_dim*n_heads] + // share the same linear element layout, so a flat add is correct. for (int64_t j = 0; j < d_model; ++j) { pos_u[j] += q_bias[j]; pos_v[j] += q_bias[j]; @@ -431,26 +400,10 @@ 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. Step 3 moved them to -// F16 in the GGUF to halve the Vulkan/Metal matmul cost, but Zen 2 -// (and anything else without native F16 compute) pays an F16→F32 -// upconvert per-element that outweighs the bandwidth win — a ~600 ms -// regression on a 9.3 s q4_k_m CPU encode. -// -// The cleanest fix without shipping two GGUFs is to hoist the -// conversion to load time: dequantize into a separate backend buffer -// and point the weight slots at the F32 copies. Vulkan/Metal/CUDA -// primary backends skip this step and keep the F16 weights. The -// original F16 tensors stay allocated in the main weight buffer -// (unused) — dropping them individually isn't supported by ggml's -// backend buffer model, and the ~235 MB cost is acceptable for the -// CPU-only path. -// -// Weight collection is the only family-specific bit: we walk -// m.weights.blocks and hand the shared helper a list of (slot, src) -// pairs. The helper does the backend check, allocation, dequantize, -// and repoint. +// 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); @@ -513,79 +466,53 @@ transcribe_status load( return st; } - // Tokenizer. if (const transcribe_status st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) { return st; } - // Hparams. if (const transcribe_status st = read_cohere_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) { return st; } - // Publish the input-length ceiling now that the encoder positional- - // encoding span and frontend rate are known (apply_family_invariants - // ran before the hparams were read). The encoder is the binding INPUT - // limit for cohere — see cohere_max_audio_ms above and - // docs/input-limits.md. + // Publish the input-length ceiling now that the encoder pos-emb span + // and frontend rate are known (the encoder is the binding INPUT limit). m->caps.max_audio_ms = cohere_max_audio_ms(m->hparams); // Basis for the session-level limits query (transcribe_session_get_limits). - // - // DISCREPANCY (intentional, flagged per docs/input-limits.md): cohere has - // TWO distinct ceilings. The INPUT bound is the encoder pos-emb table - // (enc_pos_emb_max_len, ~400 s of audio) — that is what caps.max_audio_ms - // reports. The session n_ctx knob and the KV-byte estimate, however, key off - // the DECODER self-KV ceiling (dec_max_seq, 1024): that is the value cc->n_ctx - // lowers (cohere_dec_ctx_ceiling) and the self-KV cache is sized to. Because - // the LimitsBasis is a single context cap, it is filled from the DECODER side - // so effective_n_ctx and max_kv_bytes are exact. The consequence is that the - // query's effective_max_audio_ms (= (eff - overhead - gen_reserve) * - // ms_per_audio_token) is computed off the *decoder* ceiling and so will NOT - // match caps.max_audio_ms (the encoder bound) — for cohere the decoder context - // bounds the OUTPUT transcript, not the audio input, so that derived figure is - // advisory only. caps.max_audio_ms remains the authoritative input bound. + // The LimitsBasis is a single context cap, filled from the DECODER side + // (dec_max_seq) so effective_n_ctx and max_kv_bytes are exact. Its derived + // effective_max_audio_ms therefore keys off the decoder ceiling and will + // NOT match caps.max_audio_ms (the encoder input bound) — for cohere the + // decoder context bounds the OUTPUT transcript, so that figure is advisory. if (m->hparams.dec_max_seq > 0 && m->hparams.enc_subsampling_factor > 0 && m->hparams.fe_hop_length > 0 && m->hparams.fe_sample_rate > 0) { m->limits.has_context_cap = true; // Audio bound is the encoder table (caps.max_audio_ms), not the - // decoder context, so effective_max_audio_ms must not shrink with - // n_ctx; effective_n_ctx / max_kv_bytes still reflect the decoder. + // decoder context, so effective_max_audio_ms must not shrink with n_ctx. m->limits.audio_from_caps = true; m->limits.model_max_ctx = m->hparams.dec_max_seq; - // Fixed control-token preamble that precedes the transcript in the - // decoder self-context (▁/soc/sot/emo/lang×2/pnc/noitn/notimestamp/ - // nodiarize — see run()'s prompt_pieces). Audio is in cross-KV, never in - // the self-context, so there is no audio-token overhead here. + // 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() - // ms-per-audio-token = subsampling_factor * hop_length * 1000 / sr - // (the same inversion cohere_max_audio_ms applies at T_enc). + // ms-per-audio-token = subsampling_factor * hop_length * 1000 / sr. m->limits.ms_per_audio_token = static_cast(m->hparams.enc_subsampling_factor) * m->hparams.fe_hop_length * 1000.0 / m->hparams.fe_sample_rate; - // Self-KV: the cache stores dec_hidden (n_state) per layer, two tensors - // (K and V); cohere has no GQA split, so n_state == dec_hidden. + // Self-KV stores dec_hidden per layer, two tensors (K, V); no GQA. m->limits.kv_elems_per_ctx_token = (int64_t) m->hparams.dec_hidden * m->hparams.dec_n_layers * 2; } - // Set vocab_size from tokenizer. m->hparams.vocab_size = m->tok.n_tokens(); - // Set special token IDs from tokenizer. m->hparams.bos_token_id = m->tok.bos_id(); m->hparams.eos_token_id = m->tok.eos_id(); - // Hard-fail at load time if the tokenizer did not supply an EOS - // token id. Previously this was papered over at decode time with a - // fallback to token id 3 plus a stderr warning; that is worse than - // refusing the model because (a) id 3 is a random guess that only - // happens to match this family's current checkpoint, and (b) a - // missing EOS is a GGUF-builder bug that should surface during - // conversion, not hide behind a runtime fallback in production. + // Hard-fail at load time if the tokenizer supplied no EOS token id — + // a missing EOS is a GGUF-builder bug that should surface at conversion. if (m->hparams.eos_token_id < 0) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere: GGUF tokenizer has no eos_token_id -- " @@ -606,15 +533,9 @@ transcribe_status load( cfg.f_max = m->hparams.fe_f_max; cfg.pad_mode = m->hparams.fe_pad_mode; - // Load checkpoint filterbank/window from GGUF if present. - // These small tensors are read directly from the file using - // the gguf index — avoids recomputing from scratch and - // matches the exact values the model was trained with. - // - // read_f32_tensor_checked validates type (must be F32), byte - // alignment, expected element count, and read completeness. - // Absent → compute from hparams (the non-GGUF path). - // BadType/BadSize/ReadErr → hard fail with TRANSCRIBE_ERR_GGUF. + // Load checkpoint filterbank/window from GGUF if present (exact + // trained values). read_f32_tensor_checked: Absent -> compute from + // hparams; BadType/BadSize/ReadErr -> hard fail. { using R = load_common::ReadF32Result; @@ -641,7 +562,7 @@ transcribe_status load( m->mel.emplace(cfg); } - // Stage 2: reopen with no_alloc. + // Reopen with no_alloc. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -660,13 +581,10 @@ transcribe_status load( return st; } - // Resolve the backend plan from the caller request. See - // transcribe-backend.h + transcribe-load-common.h for the - // semantics (AUTO / CPU / METAL / VULKAN). The important case - // for cohere is strict CPU: the F16→F32 conv pointwise - // promotion below depends on `plan.primary_kind == - // BackendKind::Cpu`, which only a TRANSCRIBE_BACKEND_CPU - // request reliably produces. + // Resolve the backend plan (see transcribe-load-common.h). The conv + // pointwise F16->F32 promotion below depends on a strict-CPU plan + // (primary_kind == BackendKind::Cpu), which only TRANSCRIBE_BACKEND_CPU + // reliably produces. const transcribe_backend_request backend_req = (params != nullptr) ? params->backend : TRANSCRIBE_BACKEND_AUTO; @@ -705,7 +623,6 @@ transcribe_status load( gguf_free(gguf_data); - // Fuse BatchNorm. if (const transcribe_status st = fuse_batch_norm(*m); st != TRANSCRIBE_OK) { @@ -719,10 +636,7 @@ transcribe_status load( return st; } - // On CPU backend, dequantize conv pointwise weights back to F32 — - // Zen 2 class CPUs don't have native F16 compute and the upconvert - // cost per matmul outweighs the bandwidth savings from the smaller - // weight. No-op on GPU backends. + // 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) { @@ -747,22 +661,12 @@ transcribe_status init_context( cc->model = model; cc->n_threads = params->n_threads; cc->kv_type = params->kv_type; - // Cache the caller's n_ctx knob. For cohere this lowers the DECODER - // self-KV ceiling (dec_max_seq); it does not affect the encoder input - // limit. See cohere_dec_ctx_ceiling and docs/input-limits.md. + // Cache the caller's n_ctx knob; for cohere it lowers the DECODER + // self-KV ceiling only (see cohere_dec_ctx_ceiling). cc->n_ctx = transcribe_session_params_n_ctx(params); - // Flash-attention policy -- see CohereSession (cohere.h) for why - // this is split into encoder vs decoder. The short version: - // encoder dk=160 is unsupported by upstream ggml Metal flash, so - // the encoder auto-disables on Metal; decoder dk=128 works on - // every backend we ship, so the decoder defaults on. The env-var - // overrides (TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH) are - // applied globally in transcribe::flash::apply_env_overrides. - // - // We check the classified primary BackendKind here instead of - // string-matching on ggml_backend_name — the plan already has - // the kind ready. + // Flash-attention policy — see CohereSession (cohere.h). Encoder + // auto-disables on Metal; decoder defaults on; env overrides apply globally. auto * cm = static_cast(model); const bool is_metal = (cm->plan.primary_kind == transcribe::BackendKind::Metal); @@ -793,8 +697,7 @@ transcribe_status run( return TRANSCRIBE_ERR_INVALID_ARG; } - // Pre-run abort check. Cohere ASR is single-chunk today; this is - // the single observation point. + // Pre-run abort check (cohere is single-chunk today). if (cc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } @@ -822,16 +725,11 @@ transcribe_status run( } cc->t_mel_us = ggml_time_us() - t_mel_start; - // ----- Input-length gate (see docs/input-limits.md) ------------ - // The encoder's relative positional-encoding table is the binding - // INPUT limit: encode_positions() is generated at the subsampled - // sequence length T_enc, so T_enc must stay within the trained span - // enc_pos_emb_max_len or the runtime pos table aliases past the - // trained range (a silent out-of-bounds before this gate existed). - // T_enc is a deterministic function of mel_n_frames, so reject an - // over-length clip here — before the encoder graph is even built — - // rather than paying for a compute pass that cannot fit. The - // rejection goes through the log callback, not raw stderr. + // ----- Input-length gate ----- + // T_enc must stay within the trained pos-emb span enc_pos_emb_max_len + // or the runtime pos table aliases past the trained range. T_enc is a + // deterministic function of mel_n_frames, so reject an over-length clip + // here, before building the encoder graph. if (cm->hparams.enc_pos_emb_max_len > 0) { const int t_enc_pred = cohere_predict_t_enc( mel_n_frames, cm->hparams.enc_subsampling_factor); @@ -1015,15 +913,9 @@ transcribe_status run( // const char * lang = (params && params->language) ? params->language : "en"; - // Language validation lives in the central dispatcher - // (transcribe_run), which rejects unsupported caller-provided - // languages with TRANSCRIBE_ERR_UNSUPPORTED_LANGUAGE before - // reaching this handler. The NULL → "en" default above applies - // only when the caller did not specify a language; "en" is - // always in the model's list for every published Cohere ASR - // checkpoint. If that ever ceases to hold, the dispatcher's - // check still covers non-NULL cases and we'd revisit the - // default here. + // The dispatcher (transcribe_run) rejects unsupported caller languages + // before this handler; the NULL -> "en" default applies only when the + // caller specified none. const std::string lang_token = std::string("<|") + lang + "|>"; const std::vector prompt_pieces = { @@ -1053,7 +945,7 @@ transcribe_status run( } const int prompt_len = static_cast(prompt_ids.size()); - // --- Step 1: Initialize KV cache -------------------------------- + // ----- Initialize KV cache ----- { // Free any existing cache if T_enc changed. if (cc->kv_cache.buffer != nullptr && @@ -1090,15 +982,9 @@ transcribe_status run( } } - // Helper to create a fresh compute context. - // - // Any ggml_tensor pointer we hold that was allocated inside the - // previous compute_ctx becomes dangling the instant we ggml_free - // it below, so null them out here. Today cc->encoder_out is the - // only such pointer -- its *data* was already copied to - // cc->enc_host above, and cross-attn graphs below re-declare a - // fresh encoder_out input tensor -- but keeping a stale pointer - // around invites a future edit to use-after-free it. + // Fresh compute context. Null cc->encoder_out first: it was allocated in + // the context being freed (its data is already in cc->enc_host), so the + // stale pointer must not outlive the ggml_free. auto new_compute_ctx = [&](size_t mem_size) -> bool { if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); @@ -1128,7 +1014,7 @@ transcribe_status run( return nullptr; }; - // --- Step 2: Compute cross-attention K/V ------------------------- + // ----- Compute cross-attention K/V ----- const int64_t t_dec_start = ggml_time_us(); { if (!new_compute_ctx(4 * 1024 * 1024)) { @@ -1170,7 +1056,7 @@ transcribe_status run( cc->kv_cache.cross_populated = true; } - // --- Step 3: Prompt pass with KV cache --------------------------- + // ----- Prompt pass with KV cache ----- { if (!new_compute_ctx(4 * 1024 * 1024)) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -1178,9 +1064,8 @@ transcribe_status run( return TRANSCRIBE_ERR_OOM; } - // Prompt pass: skip log_softmax and emit a GPU argmax over the - // last position. Debug dumps still need the pre-head hidden - // state; when dumping is enabled we take the slower path. + // Skip log_softmax and emit a GPU argmax over the last position. + // Debug dumps need the pre-head state, so they take the slower path. const bool prompt_skip_softmax = !transcribe::debug::enabled(); DecoderBuild db = build_decoder_graph_kv( cc->compute_ctx, cm->weights, cm->hparams, @@ -1332,25 +1217,15 @@ transcribe_status run( cc->has_result = true; }; - // --- Step 4: Autoregressive step passes ---------------------- - // - // Two decoder loop variants; pick by primary backend kind. - // - // GPU (Vulkan/Metal/CUDA/SYCL): build_step_graph — one static- - // topology graph for the whole utterance. KV writes go via - // ggml_set_rows at runtime kv_idx; flash-attn reads a fixed - // max_n_kv window with a runtime mask. Removes per-step - // graph_build + sched_alloc, which dominate dispatch overhead - // on GPUs. - // - // CPU (incl. accelerator host-memory backends): build_decoder_ - // graph_kv per step — n_kv grows with n_past, attention only - // reads the populated prefix. CPU has no dispatch overhead - // to amortize, so the static-graph bandwidth tax (reading - // max_n_kv KV slots when avg n_past is small) is a net loss. - // - // The prompt pass uses build_decoder_graph_kv on both paths - // (already executed above) — only the per-token loop branches. + // ----- Autoregressive step passes ----- + // Two loop variants by primary backend kind: + // GPU: build_step_graph — one static-topology graph; KV writes via + // ggml_set_rows at runtime kv_idx, flash-attn over a fixed + // max_n_kv window. Avoids per-step graph_build + sched_alloc, + // which dominate GPU dispatch overhead. + // CPU: build_decoder_graph_kv per step — n_kv grows with n_past + // (reads only the populated prefix). CPU has no dispatch overhead + // to amortize, so the static-graph bandwidth tax is a net loss. int n_past = prompt_len; const bool primary_is_gpu = @@ -1360,10 +1235,9 @@ transcribe_status run( if (primary_is_gpu) { // ---------- Static-graph step path (GPU) ---------- - // max_n_kv: pad to next power of two with a 1024 floor. - // Vulkan/Metal flash-attn dispatches faster on pow2 ne[1]; - // the slight bandwidth cost is amortized by removing - // per-step graph_build + sched_alloc. + // max_n_kv: next power of two (1024 floor) — Vulkan/Metal flash + // dispatches faster on pow2 ne[1], and the static graph amortizes + // the slight bandwidth cost. int max_n_kv = 1024; while (max_n_kv < prompt_len + max_tokens) max_n_kv *= 2; if (max_n_kv > cc->kv_cache.n_ctx) max_n_kv = cc->kv_cache.n_ctx; @@ -1408,9 +1282,8 @@ transcribe_status run( return TRANSCRIBE_ERR_ABORTED; } if (n_past + 1 > max_n_kv) { - // Output ran into the self-KV window before EOS. Keep - // the partial transcript, flag truncation, warn — never - // discard a usable result. See docs/input-limits.md. + // Hit the self-KV window before EOS: keep the partial + // transcript, flag truncation, warn — never discard it. cc->was_truncated = true; transcribe::log_msg( TRANSCRIBE_LOG_LEVEL_WARN, @@ -1456,9 +1329,8 @@ transcribe_status run( } } else { // ---------- Dynamic-graph step path (CPU) ---------- - // Reserve scheduler buffers with a worst-case single-token - // graph (maximum n_past = n_ctx - 1). This ensures that - // alloc_graph during the loop never triggers reallocation. + // Reserve scheduler buffers with a worst-case single-token graph + // (n_past = n_ctx - 1) so alloc_graph never reallocates in-loop. if (new_compute_ctx(4 * 1024 * 1024)) { const int worst_n_past = cc->kv_cache.n_ctx - 1; DecoderBuild db_reserve = build_decoder_graph_kv( @@ -1481,9 +1353,8 @@ transcribe_status run( } if (n_past + 1 > cc->kv_cache.n_ctx) { - // Output filled the self-KV cache before EOS. Keep the - // partial transcript, flag truncation, warn — never - // discard a usable result. See docs/input-limits.md. + // Filled the self-KV cache before EOS: keep the partial + // transcript, flag truncation, warn — never discard it. cc->was_truncated = true; transcribe::log_msg( TRANSCRIBE_LOG_LEVEL_WARN, @@ -1538,14 +1409,10 @@ transcribe_status run( } } - // A complete decode stops because the model emitted EOS - // (next_token == eos_id). If instead the loop ran out of budget — - // it reached the 512 max-new cap (step == max_tokens) with a - // non-EOS token still pending — the transcript is truncated. The - // two break sites above already flag the self-KV-full case; this - // mirror covers the normal-exit-at-budget case. A clean EOS decode - // leaves next_token == eos_id and never sets the flag. See - // docs/input-limits.md. + // A clean decode stops at EOS. If the loop instead ran out of budget + // (max-new cap) with a non-EOS token pending, the transcript is + // truncated. The break sites flag the self-KV-full case; this covers + // the normal-exit-at-budget case. if (!cc->was_truncated && next_token != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -1556,35 +1423,15 @@ transcribe_status run( static_cast(generated_ids.size())); } - // Build result hierarchy. Cohere advertises - // max_timestamp_kind == NONE: "text but no alignment data". - // That means no per-token, per-word, or per-segment timing - // surface at all — not even a zero-timed placeholder. The - // honest output shape is: - // - // full_text - the whole transcript - // segments[0] - single segment whose text == full_text, - // with zeroed timings and zero first/n - // counts (the caller asked for NONE, so - // they should not see a token or word - // structure) - // tokens - empty - // words - empty - // - // The previous implementation populated cc->tokens with - // zero-timed entries and set seg.n_tokens to the token - // count, which made transcribe_n_tokens > 0 and let a - // caller enumerate token_text / token_t0_ms for a model - // that is not supposed to expose alignment data. That - // undercut the NONE contract. Drop them. + // Build the result. max_timestamp_kind == NONE means text but no + // alignment data: full_text plus one segment (text == full_text, + // zeroed timings and counts), with empty tokens/words. commit_result(); - } // end Step 3 block + } - // Output truncation is a hard status: when the single-shot decode stopped - // before EOS (either break site above, or the post-loop budget check) it - // set cc->was_truncated. The partial transcript is already committed and - // stays readable (like an aborted run); surface the truncation to the - // caller rather than reporting a clean OK. See docs/input-limits.md. + // Output truncation is a hard status: the partial transcript is committed + // and stays readable (like an aborted run), but we surface the truncation + // rather than reporting a clean OK. return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } @@ -1592,18 +1439,15 @@ transcribe_status run( // Offline batched decode (transcribe_run_batch) // =========================================================================== // -// The encoder is compute-bound and stays SERIAL per utterance (the granite -// lesson: batching a heavy conformer is a wash and regresses on mixed -// lengths). Only the host-side mel is parallelized and the autoregressive -// DECODE is batched — decode at one-token-per-step is memory-bandwidth- -// bound, so batching B utterances amortizes the per-step weight reads. +// The compute-bound encoder stays SERIAL per utterance (batching a heavy +// conformer regresses on mixed lengths). The host-side mel is parallelized +// and the bandwidth-bound autoregressive DECODE is batched, so B utterances +// amortize the per-step weight reads. // -// Cross-attention is the cohere-specific wrinkle vs the causal_lm families: -// each utterance has its own encoder output (variable T_enc), so the cross -// KV cache carries a batch dim and a per-utterance cross-pad mask discards -// the frames past T_enc[b]. The short, uniform prompt is fed through the -// same batched step graph (prompt_len sequential steps) rather than a -// separate prompt graph. +// Cohere-specific wrinkle: each utterance has its own encoder output +// (variable T_enc), so the cross KV cache carries a batch dim and a +// per-utterance cross-pad mask discards frames past T_enc[b]. The uniform +// prompt feeds through the same batched step graph (prompt_len steps). // Encoder for one utterance from a precomputed mel buffer → host [hidden, T_enc]. transcribe_status encode_one_to_host( @@ -1772,10 +1616,9 @@ transcribe_status run_batch( mel_us += ggml_time_us() - t_mel0; // ----- Pass 1: serial per-utterance encoder ----- - // Per-utterance input-length gate (same binding limit as run(): the encoder - // positional-encoding span enc_pos_emb_max_len). An over-length utterance is - // marked invalid with INPUT_TOO_LONG so the rest of the batch still decodes, - // mirroring how an encode failure drops one row. See docs/input-limits.md. + // Per-utterance input-length gate (same enc_pos_emb_max_len limit as + // run()). An over-length utterance is marked invalid with INPUT_TOO_LONG + // so the rest of the batch still decodes. std::vector reject_status(n, TRANSCRIBE_ERR_INVALID_ARG); std::vector> enc_hosts(n); std::vector T_enc(n, 0); @@ -1884,12 +1727,10 @@ transcribe_status run_batch( } // ----- Batched step graph with a GROWING self-attention window ----- - // The self-KV holds only the short text prompt + transcript (audio is in - // cross-KV), so the static n_ctx-wide read is mostly empty for short - // clips. Start the read window at 64 and double it as n_past advances - // (rebuild graph + widen mask, O(log) rebuilds; the KV cache persists), - // keeping per-step self-KV bandwidth within 2× of the real n_past. The - // cache capacity (max_n_kv) is unchanged. + // Self-KV holds only the short prompt + transcript (audio is in cross-KV), + // so a static n_ctx-wide read is mostly empty for short clips. Start the + // read window at 64 and double it as n_past advances (O(log) rebuilds; the + // KV cache persists), keeping per-step self-KV bandwidth within 2x of real. const ggml_fp16_t f16_zero = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t f16_ninf = ggml_fp32_to_fp16(-INFINITY); const int32_t eos_id = hp.eos_token_id; @@ -1952,10 +1793,8 @@ transcribe_status run_batch( rs.full_text = full; rs.result_kind = TRANSCRIBE_TIMESTAMPS_NONE; rs.has_result = true; rs.status = TRANSCRIBE_OK; - // Per-utterance truncation parity with the single-shot path: a valid row - // that hit the generation budget / context window before eos reports - // TRANSCRIBE_ERR_OUTPUT_TRUNCATED (partial transcript retained). Only - // override an otherwise-OK status — never a worse one. + // Truncation parity with the single-shot path; only override an + // otherwise-OK status, never a worse one. if (rs.status == TRANSCRIBE_OK && b < static_cast(truncated.size()) && truncated[b]) { cc->was_truncated = true; diff --git a/src/arch/cohere/weights.cpp b/src/arch/cohere/weights.cpp index 27ee6be6..510c5fc5 100644 --- a/src/arch/cohere/weights.cpp +++ b/src/arch/cohere/weights.cpp @@ -22,10 +22,6 @@ namespace transcribe::cohere { -// --------------------------------------------------------------------------- -// Hparams -// --------------------------------------------------------------------------- - namespace { constexpr const char * kFamilyTag = "cohere"; @@ -170,20 +166,12 @@ transcribe_status read_cohere_hparams(const gguf_context * gguf, return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Weights -// --------------------------------------------------------------------------- - namespace { using transcribe::weights::lname; -// The canonical find_tensor() + lname() helpers live in -// src/transcribe-weights-util.{h,cpp}; see that header for rationale. -// They are shared between every per-family weights.cpp. The GET_* -// macros still live here so the family log tag lands in diagnostics, -// but the type allowlists (TRANSCRIBE_QUANT_{LINEAR,CONV}_TYPES) -// are project-wide — every family accepts the same set. +// find_tensor() + lname() live in src/transcribe-weights-util.{h,cpp}; +// the GET_* macros stay here so the family log tag lands in diagnostics. constexpr const char * kTag = kFamilyTag; #define GET_F32(slot, name, ...) \ @@ -316,15 +304,14 @@ transcribe_status build_cohere_weights(ggml_context * ctx_meta, GET_F32(weights.enc_dec_proj.bias, "enc_dec_proj.bias", dec_h); // ----- decoder embedding ----- - // Token embedding: ne=[dec_hidden, vocab_size] in ggml (PyTorch [vocab, hidden]). - // We read vocab_size from the tensor shape since we don't have it as a separate KV. + // Token embedding ne=[dec_hidden, vocab_size]; vocab_size is read from + // the tensor shape (no separate KV). { ggml_tensor * tw = ggml_get_tensor(ctx_meta, "dec.embed.token.weight"); if (tw == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere: missing tensor \"dec.embed.token.weight\""); return TRANSCRIBE_ERR_GGUF; } - // The embedding table: ne[0]=dec_hidden, ne[1]=vocab_size. if (tw->ne[0] != dec_h) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "cohere: dec.embed.token.weight ne[0]=%lld, expected %lld", @@ -335,10 +322,8 @@ transcribe_status build_cohere_weights(ggml_context * ctx_meta, weights.dec_embed.token_w = tw; } - // Read vocab_size from the token embedding shape. const int64_t vocab_size = weights.dec_embed.token_w->ne[1]; - // Positional encoding: [dec_hidden, dec_max_seq]. GET_F32(weights.dec_embed.pos_enc, "dec.embed.pos_enc", dec_h, hp.dec_max_seq); GET_F32(weights.dec_embed.norm_w, "dec.embed.norm.weight", dec_h); GET_F32(weights.dec_embed.norm_b, "dec.embed.norm.bias", dec_h); diff --git a/src/arch/cohere/weights.h b/src/arch/cohere/weights.h index a6bbdd07..0f352160 100644 --- a/src/arch/cohere/weights.h +++ b/src/arch/cohere/weights.h @@ -32,10 +32,6 @@ struct ggml_tensor; namespace transcribe::cohere { -// --------------------------------------------------------------------------- -// Hyperparameters -// --------------------------------------------------------------------------- - struct CohereHParams { // Encoder (Conformer). int32_t enc_n_layers = 0; @@ -90,10 +86,6 @@ struct CohereHParams { transcribe_status read_cohere_hparams(const gguf_context * gguf, CohereHParams & hp); -// --------------------------------------------------------------------------- -// Weight slots -// --------------------------------------------------------------------------- - struct CoherePreEncode { // Same structure as Parakeet: dw_striding subsampling. ggml_tensor * conv0_w = nullptr; diff --git a/src/arch/funasr_nano/adaptor.cpp b/src/arch/funasr_nano/adaptor.cpp index 0dce8f94..8075e3c7 100644 --- a/src/arch/funasr_nano/adaptor.cpp +++ b/src/arch/funasr_nano/adaptor.cpp @@ -25,8 +25,7 @@ ggml_tensor * named(ggml_tensor * t, const char * name) { // Weight matmul forcing F32 accumulation for F16 weights, so CUDA's cuBLAS // path does not accumulate the adaptor's multi-token GEMMs in F16 (which // overflows F16's ~65504 range -> NaNs). Gated on F16 so BF16/quantized/F32 -// weights and CPU/Metal are bit-identical to before. See the matching helper -// in sanm.cpp / causal_lm.cpp for the full rationale. +// weights and CPU/Metal are bit-identical. See sanm.cpp / causal_lm.cpp. ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * x) { ggml_tensor * y = ggml_mul_mat(ctx, w, x); if (w->type == GGML_TYPE_F16) { @@ -166,7 +165,6 @@ AdaptorBuild build_adaptor_graph(ggml_context * ctx, x = ggml_add(ctx, x, w.adaptor.linear1_b); mark_dump(ab.dumps.linear1_out, x, "adaptor.linear1.out"); - // ReLU. x = ggml_relu(ctx, x); // linear2 (2048 → 1024) + bias. diff --git a/src/arch/funasr_nano/adaptor.h b/src/arch/funasr_nano/adaptor.h index 8a954442..19056bdc 100644 --- a/src/arch/funasr_nano/adaptor.h +++ b/src/arch/funasr_nano/adaptor.h @@ -1,14 +1,9 @@ // arch/funasr_nano/adaptor.h - audio adaptor graph builder. // // 2-layer transformer adaptor (FunASR's "Transformer" adaptor class): -// -// x_in [encoder_dim=512, T_lfr] -// -> linear1 [encoder_dim, pre_ffn_dim=2048] + bias = adaptor.linear1.out -// -> ReLU -// -> linear2 [pre_ffn_dim, llm_dim=1024] + bias = adaptor.linear2.out -// -> blocks[0..1]: pre-LN MHA (LayerNorm eps=1e-12) + bottleneck FFN -// (1024 → 256 → 1024, ReLU, biases) = adaptor.blocks.0.out -// adaptor.out +// x_in [encoder_dim=512, T_lfr] -> linear1 (512->2048) -> ReLU -> +// linear2 (2048->llm_dim=1024) -> blocks[0..1] (pre-LN MHA + bottleneck FFN +// 1024->256->1024) -> [llm_dim, T_lfr]. // // downsample_rate=1 means linear1 input is encoder_dim (no fold). Only // retained as a configuration knob for sibling variants that may set k>1. diff --git a/src/arch/funasr_nano/capabilities.cpp b/src/arch/funasr_nano/capabilities.cpp index 0ed1910c..e02599c9 100644 --- a/src/arch/funasr_nano/capabilities.cpp +++ b/src/arch/funasr_nano/capabilities.cpp @@ -22,11 +22,6 @@ void apply_family_invariants(transcribe_model & model) { // ("，不进行文本规整" appended on itn=false, omitted on itn=true). // The generic transcribe_run_params::itn enum routes here. No PNC // runtime toggle. - // - // TODO(family doc): observe whether itn=true bundles - // punctuation/casing changes alongside number/date normalization on - // shipped variants (fun-asr-nano-2512, fun-asr-mlt-nano-2512) and - // record in the family doc. API shape unchanged either way. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); transcribe::set_feature(&model, TRANSCRIBE_FEATURE_ITN, true); } diff --git a/src/arch/funasr_nano/decoder.cpp b/src/arch/funasr_nano/decoder.cpp index d0a56392..899a480d 100644 --- a/src/arch/funasr_nano/decoder.cpp +++ b/src/arch/funasr_nano/decoder.cpp @@ -6,7 +6,7 @@ // - graph allocation // - audio injection (3-way concat: prefix | adaptor_out | suffix); // T_audio == 0 is a supported chat-only path -// - tensor naming and dump-point preservation for validate.py parity +// - tensor naming and dump points // - final RMSNorm + tied lm_head head // - block_prefill / block_step driver loops // @@ -14,12 +14,9 @@ // - The audio block is the adaptor output (already in llm_dim=1024 // space), not the raw encoder output. // - T_audio is allowed to be 0 (no audio); the prefill skips the -// concat and uses token_emb_all directly. Kept here at the call -// site since the shared block helper has no business knowing about -// audio shapes. +// concat and uses token_emb_all directly. // - The step graph also exposes raw logits as an output (not just -// argmax) so the autoregressive driver can dump `dec.logits_raw.gen8` -// for tensor parity at gen step 7. +// argmax) so the autoregressive driver can dump them at gen step 7. #include "decoder.h" diff --git a/src/arch/funasr_nano/decoder.h b/src/arch/funasr_nano/decoder.h index b7ece5d5..95f52ac8 100644 --- a/src/arch/funasr_nano/decoder.h +++ b/src/arch/funasr_nano/decoder.h @@ -60,10 +60,7 @@ struct StepBuild { ggml_tensor * kv_idx_in = nullptr; // [1] i64 ggml_tensor * mask_in = nullptr; // [max_n_kv, 1] f16 ggml_tensor * out = nullptr; // [1] i32 — argmax - // Pre-argmax logits, exposed for tensor-validation dumps. The - // graph keeps both alive; reading `logits` requires - // ggml_backend_tensor_get since it lives on the device. - ggml_tensor * logits = nullptr; // [vocab] f32 + ggml_tensor * logits = nullptr; // [vocab] f32 — pre-argmax, for dumps ggml_cgraph * graph = nullptr; int max_n_kv = 0; diff --git a/src/arch/funasr_nano/encoder.cpp b/src/arch/funasr_nano/encoder.cpp index e27f7ae1..4fa8940c 100644 --- a/src/arch/funasr_nano/encoder.cpp +++ b/src/arch/funasr_nano/encoder.cpp @@ -1,19 +1,9 @@ // arch/funasr_nano/encoder.cpp - SAN-M encoder graph (no prefix prepend, // no CTC head). // -// Forward shape (single-utterance, batch=1): -// -// frontend.in [d_input=560, T_lfr] -// -> scale by sqrt(d_model=512) -// -> add sinusoidal PE (depth = d_input, 1-based positions) -// = enc.embed.out -// -> encoders0[0] (SAN-M block, projection 560 -> 512, no attn-residual) -// = enc.encoders0.0.out -// -> encoders[0..n-2] x49 (residual SAN-M blocks) -// = enc.encoders.{0,24,48}.out -// -> after_norm (LayerNorm eps=1e-12) = enc.after_norm.out -// -> tp_encoders[0..tp-1] x20 (residual SAN-M) = enc.tp_encoders.{0,10,19}.out -// -> tp_norm (LayerNorm eps=1e-12) = enc.tp_norm.out +// frontend.in [d_input=560, T_lfr] -> scale by sqrt(d_model=512) + sinusoidal +// PE -> encoders0[0] (SAN-M projection 560->512) -> encoders x49 (residual +// SAN-M) -> after_norm -> tp_encoders x20 -> tp_norm [d_model, T_lfr]. // // SAN-M block topology lives in src/sanm/sanm.h (shared with sensevoice). // This file owns the surrounding shape — embedding scale + PE add, diff --git a/src/arch/funasr_nano/model.cpp b/src/arch/funasr_nano/model.cpp index a802f1b8..985e5ddf 100644 --- a/src/arch/funasr_nano/model.cpp +++ b/src/arch/funasr_nano/model.cpp @@ -87,22 +87,17 @@ constexpr const char k_default_variant[] = "fun-asr-nano-2512"; // Input-length contract (see docs/input-limits.md). Fun-ASR-Nano is a // hard-context-cap family: audio tokens + chat prompt + generation share the // Qwen3 decoder's context window (dec_max_position_embeddings). The KV cache -// grows to fit the prompt, clamped to that ceiling (replacing the earlier -// fixed 2048 wall that ignored the model's real context). Over-length input -// is rejected up front with TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that +// grows to fit the prompt, clamped to that ceiling. Over-length input is +// rejected up front with TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that // fills the per-run generation budget before end-of-stream is flagged via // transcribe_was_truncated(). // --------------------------------------------------------------------------- -// Per-run generation budget (matches the step-loop literal below; the decode -// loop is intentionally left unchanged — see the family doc). +// Per-run generation budget. constexpr int k_max_new = 256; // Effective decoder context ceiling, in tokens: the model's trained maximum, -// optionally lowered — never raised — by the caller's session n_ctx knob. For -// Fun-ASR-Nano dec_max_position_embeddings is the Qwen3 RoPE max (40960), huge -// relative to any practical clip; grow-to-fit means we never allocate it all -// at once, so using it as the ceiling matches the qwen3_asr reference exactly. +// optionally lowered — never raised — by the caller's session n_ctx knob. int funasr_nano_context_ceiling(int32_t n_ctx_knob, const FunAsrNanoHParams & hp) { int ceiling = hp.dec_max_position_embeddings; if (n_ctx_knob > 0 && n_ctx_knob < ceiling) { @@ -113,19 +108,9 @@ int funasr_nano_context_ceiling(int32_t n_ctx_knob, const FunAsrNanoHParams & hp // Advisory transcribe_capabilities::max_audio_ms: the longest audio whose // audio tokens plus a representative prompt and the generation reserve fit -// the context ceiling. The audio-token count is a deterministic function of -// sample count: -// T_frames = 1 + floor((N - win) / hop) ~ N/hop -// T_lfr = ceil(T_frames / lfr_n) ~ T_frames/lfr_n -// audio_tokens = fake_token_len(T_lfr) (3 stride-2 folds) ~ T_lfr/8 -// so audio_tokens ≈ N / (hop * lfr_n * 8). Inverting: -// ms ≈ audio_tokens * 8 * lfr_n * hop * 1000 / sr. -// Returns 0 ("unknown / unbounded") if the rate constants are missing. Note: -// the ceiling (40960) is huge, so the number is large and honest — within it, -// a long transcript can still truncate at the generation budget -// (transcribe_was_truncated). We report it as-is (no 1-hour cap); a fully -// dense clip near the ceiling is not physically reachable from real audio, so -// capping the report would only hide the honest derived value. +// the context ceiling. audio_tokens ≈ N / (hop * lfr_n * folds), so +// ms ≈ audio_tokens * folds * lfr_n * hop * 1000 / sr. Returns 0 +// ("unknown / unbounded") if the rate constants are missing. int64_t funasr_nano_max_audio_ms(const FunAsrNanoHParams & hp) { if (hp.dec_max_position_embeddings <= 0 || hp.fe_hop_length <= 0 || hp.fe_sample_rate <= 0 || @@ -138,8 +123,6 @@ int64_t funasr_nano_max_audio_ms(const FunAsrNanoHParams & hp) { if (max_audio_tokens <= 0) { return 0; } - // audio_tokens ≈ T_lfr / 8 ; T_lfr ≈ T_frames / lfr_n ; T_frames ≈ N / hop - // => N ≈ audio_tokens * 8 * lfr_n * hop ; ms = N * 1000 / sr. const int folds = hp.adaptor_use_low_frame_rate ? 8 : 1; const int64_t samples = static_cast(max_audio_tokens) * folds * hp.fe_lfr_n * hp.fe_hop_length; @@ -340,8 +323,6 @@ transcribe_status load( m->limits.model_max_ctx = m->hparams.dec_max_position_embeddings; m->limits.prompt_overhead = 48; m->limits.gen_reserve = k_max_new; - // audio_tokens ≈ N / (hop * lfr_n * folds) ; ms = N * 1000 / sr - // => ms per audio token = folds * lfr_n * hop * 1000 / sr m->limits.ms_per_audio_token = static_cast(folds) * m->hparams.fe_lfr_n * m->hparams.fe_hop_length * 1000.0 / m->hparams.fe_sample_rate; @@ -366,7 +347,7 @@ transcribe_status load( return TRANSCRIBE_ERR_GGUF; } - // Stage 2: reopen with no_alloc to bind the tensor catalog. + // Reopen with no_alloc to bind the tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -705,8 +686,9 @@ transcribe_status run( // Generic transcribe_run_params::itn routes here. DEFAULT maps to the // family's shipping default (use_itn=false; matches the upstream // `itn=False` Python path). OFF / ON override explicitly. The - // dispatcher's advisory WARN only fires when supports_itn == false; - // funasr_nano advertises supports_itn = true, so no WARN here. + // dispatcher's advisory WARN only fires when transcribe_model_supports( + // model, TRANSCRIBE_FEATURE_ITN) is false; funasr_nano sets + // TRANSCRIBE_FEATURE_ITN so the probe returns true and no WARN here. bool use_itn = false; if (params != nullptr) { switch (params->itn) { @@ -752,9 +734,7 @@ transcribe_status run( // Size to hold the prompt plus the generation budget, rounded up to a // 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. For typical - // short clips this is the same 1024/2048 width as the prior fixed cache, - // so the decode is unchanged — only previously-walled long clips now fit. + // a pre-allocated smaller cache is freed and re-allocated. int want_n_ctx = 1024; while (want_n_ctx < T_prompt + k_max_new) want_n_ctx *= 2; if (want_n_ctx > ceiling) want_n_ctx = ceiling; @@ -898,13 +878,12 @@ transcribe_status run( // ---- Step loop ---- const int32_t eos_id = hp.eos_token_id; - const int max_new = k_max_new; // matches the per-run generation budget + const int max_new = k_max_new; int cur_past = T_prompt; // Static step-graph shape: T_prompt prefilled + up to max_new generated. // Clamp to cc->kv_cache.n_ctx (the grown cache size) so the attention - // width never exceeds the allocation. For typical short clips this is the - // same 1024/2048 width as before — decode numerics are unchanged. + // width never exceeds the allocation. int max_n_kv = 1024; while (max_n_kv < T_prompt + max_new) max_n_kv *= 2; if (max_n_kv > cc->kv_cache.n_ctx) max_n_kv = cc->kv_cache.n_ctx; @@ -938,16 +917,8 @@ transcribe_status run( const ggml_fp16_t mask_neg_inf = ggml_fp32_to_fp16(-INFINITY); std::vector step_mask(max_n_kv, mask_neg_inf); - // Mid-generation logits dump. The reference dumper captures - // `lm_head_h.values[8]` — the 9th lm_head call (0-indexed, where - // call 0 is the prefill pass). In our step loop, the 9th lm_head - // call corresponds to iter 7 of the loop (prefill = 1st call, - // iter K = (K+2)th call), so we dump when n_steps == 7. Misnamed - // historically as gen_dump_step=8 — that captured the 10th call - // and produced an off-by-one against the reference; the bug - // happened to pass nano's tolerances because nano's adjacent-step - // logits are similar in magnitude, but blew up on MLT where the - // tied-lm_head distribution is heavily shifted. + // Mid-generation logits dump. The reference captures the 9th lm_head call + // (prefill = 1st call, iter K = (K+2)th call), so dump when n_steps == 7. const int gen_dump_step = 7; int n_steps = 0; while (next_tok != eos_id && @@ -985,12 +956,6 @@ transcribe_status run( next_tok = argmax_tok; generated_ids.push_back(next_tok); - // Mid-generation tensor coverage: `dec.logits_raw.gen8` is the - // logits for the 9th lm_head call (= step 8 of our step loop, - // when we've already emitted prefill + 7 generated tokens and - // are emitting the 8th from gen step 8). The reference dumper - // calls this same tensor `dec.logits_raw.gen8`; capture here so - // validate.py can compare the n_past>0 step-graph code path. if (n_steps == gen_dump_step && transcribe::debug::enabled()) { try_dump("dec.logits_raw.gen8", sb.logits, "decoder.logits.gen8"); } diff --git a/src/arch/funasr_nano/weights.cpp b/src/arch/funasr_nano/weights.cpp index 30992608..b0744189 100644 --- a/src/arch/funasr_nano/weights.cpp +++ b/src/arch/funasr_nano/weights.cpp @@ -1,9 +1,5 @@ // arch/funasr_nano/weights.cpp - read_funasr_nano_hparams + // build_funasr_nano_weights. -// -// Pattern mirrors src/arch/qwen3_asr/weights.cpp (audio-llm) and -// src/arch/sensevoice/weights.cpp (the encoder catalog is a pruned fork -// of sensevoice — no prefix-embedding table, no CMVN, no CTC head). #include "weights.h" diff --git a/src/arch/gigaam/decoder.cpp b/src/arch/gigaam/decoder.cpp index 53527354..6412a062 100644 --- a/src/arch/gigaam/decoder.cpp +++ b/src/arch/gigaam/decoder.cpp @@ -1,16 +1,11 @@ -// arch/gigaam/decoder.cpp - GigaAM RNN-T greedy decode on host. +// arch/gigaam/decoder.cpp - GigaAM RNN-T / CTC greedy decode on host. // -// The decoder is tiny (1 LSTM layer of 320, single joint linear) so the -// per-step compute fits cleanly in CPU cache; running it on host avoids -// backend dispatch overhead for the inner symbol-loop. Mirrors -// parakeet's host-decode pattern. -// -// PyTorch nn.LSTM stores its gate weights as a single concatenated -// [4*hidden, hidden] tensor in (i, f, g, o) gate order. The converter -// emits Wx and Wh in that layout (ggml ne=[hidden, 4*hidden]) and -// collapses bias_ih + bias_hh into a single [4*hidden] bias. The host -// LSTM step here reads them in numpy-row-major [4*hidden, hidden] -// orientation, which matches the byte layout directly. +// The decoder is tiny (1 LSTM layer of 320, single joint linear) so per-step +// compute fits in CPU cache; running on host avoids backend dispatch overhead +// for the inner symbol-loop. PyTorch nn.LSTM stores gate weights concatenated +// [4*hidden, hidden] in (i, f, g, o) order; the converter emits Wx/Wh in that +// layout and collapses bias_ih + bias_hh into one [4*hidden] bias, which the +// host LSTM step reads directly. #include "decoder.h" @@ -40,15 +35,10 @@ namespace transcribe::gigaam { namespace { -// Copy a ggml tensor's data into a host f32 vector, dequantizing if -// needed. F32 takes the fast path (single backend_tensor_get). F16, -// BF16, and every Q* / IQ* preset register a to_float in ggml's -// type-traits table; we stage the raw bytes off the backend and walk -// to_float to produce fp32. Mirrors parakeet's read_tensor_to_f32. -// -// The host RNN-T / CTC greedy decoders compute in fp32 only, so this -// runs exactly once at load time per quantized GGUF and the per-step -// hot path pays nothing. +// Copy a ggml tensor's data into a host f32 vector, dequantizing if needed. +// F32 takes the fast path; F16/BF16/Q*/IQ* stage raw bytes off the backend and +// walk ggml's to_float trait. Runs once at load time, so the per-step hot path +// pays nothing. void readback_f32(const ggml_tensor * t, std::vector & out) { if (t == nullptr) { out.clear(); @@ -164,7 +154,6 @@ static void lstm_step(const float * x, scratch_gates[i] = acc; } - // Split gates: i = 0..H, f = H..2H, g = 2H..3H, o = 3H..4H. const float * gi = scratch_gates; const float * gf = scratch_gates + H; const float * gg = scratch_gates + 2 * H; @@ -388,8 +377,7 @@ transcribe_status decode_ctc_greedy(const HostDecoderWeights & host, frame, n_classes, d_model, row); } - // Dump raw logits before log_softmax (matches reference dump - // `ctc.logits.raw`). Shape: [T_enc, n_classes] row-major. + // Raw logits before log_softmax, [T_enc, n_classes] row-major. if (transcribe::debug::enabled()) { const long long shape[2] = { T_enc, n_classes }; transcribe::debug::dump_host_f32( diff --git a/src/arch/gigaam/decoder.h b/src/arch/gigaam/decoder.h index 8057c5ca..29b7d0e2 100644 --- a/src/arch/gigaam/decoder.h +++ b/src/arch/gigaam/decoder.h @@ -1,8 +1,4 @@ // arch/gigaam/decoder.h - GigaAM RNN-T / CTC greedy decoders. -// -// M1 stub. M3 lands the RNN-T predictor LSTM + joint + greedy loop on -// host (mirrors parakeet's host-decode pattern). M4 lands the CTC -// argmax-collapse path. #pragma once diff --git a/src/arch/gigaam/encoder.cpp b/src/arch/gigaam/encoder.cpp index 7a8d6b0f..529d6b05 100644 --- a/src/arch/gigaam/encoder.cpp +++ b/src/arch/gigaam/encoder.cpp @@ -1,19 +1,9 @@ // arch/gigaam/encoder.cpp - GigaAM Conformer encoder graph builder. // -// Reuses transcribe::conformer helpers for the macaron FF, conv module, -// and the LayerNorm primitive. Adds gigaam-specific pieces: -// -// - 2-conv1d pre_encode (not parakeet's 4-conv2d stack). -// - Rotary self-attention with PRE-projection rotation (the rotation -// is applied to x before Wq/Wk; same rotated x feeds both, Wv gets -// unrotated x). -// -// Layout convention: -// - mel_in ne=[T_mel, n_mels, 1, 1] f32 (matches the reference dumper's -// [n_mels, T_mel] row-major layout byte-for-byte). -// - After pre_encode: [d_model, T_enc, 1, 1] where T_enc = floor((T_mel-1)/4) -// (two stride-2 conv1d with kernel=5, sym pad=2). -// - All conformer block intermediates stay at [d_model, T_enc, 1, 1]. +// Reuses transcribe::conformer helpers (macaron FF, conv module, LayerNorm) +// plus gigaam-specific pieces: a 2-conv1d pre_encode and rotary self-attention +// with PRE-projection rotation (rotation applied to x before Wq/Wk; Wv gets +// unrotated x). Activations stay [d_model, T_enc, B]; T_enc = floor((T_mel-1)/4). #include "encoder.h" @@ -65,32 +55,15 @@ struct PreEncodeMasks { }; // Build the 2-conv1d pre-encode stack. -// -// Input mel_in: ne=[T_mel, n_mels, B] (B == 1 single-shot) -// Output pre_enc: ne=[d_model, T_enc, B] with T_enc ≈ T_mel / 4 -// -// Reference (gigaam.encoder.StridingSubsampling): -// x = audio.transpose(1, 2) # [B, n_mels, T] -// x = self.conv(x) # Sequential(Conv1d(64→768, k=5, s=2), -// # ReLU, -// # Conv1d(768→768, k=5, s=2), -// # ReLU) -// x = x.transpose(1, 2) # [B, T_enc, d_model] -// -// In ggml: mel_in is already [T_mel, n_mels] (matching PyTorch [n_mels, T] when -// considering row-major + ggml's fast-to-slow convention). conv_1d_f32 expects -// data ne=[W=T, IC, N] and kernel ne=[K, IC, OC] — both match the GGUF layout, -// and conv_1d_f32 carries the batch N at ne[2] natively. After two stride-2 -// conv1d the result is [T_enc, d_model, B] which we transpose+contiguous to -// [d_model, T_enc, B] to match the conformer downstream layout. +// Input mel_in: ne=[T_mel, n_mels, B] (B == 1 single-shot) +// Output pre_enc: ne=[d_model, T_enc, B] with T_enc ≈ T_mel / 4 // // masks (optional): variable-length batching zero-pads each utterance's mel to -// the batch T_max along time. The convs carry bias + ReLU, so a padded zero -// becomes non-zero after the first conv and then leaks into each utterance's -// last VALID frame via the next stride-2 conv's receptive field — corrupting -// the boundary frame. NeMo's masked subsampling fix is to zero each conv -// intermediate's padded time region after every ReLU; passing non-null masks -// applies that (mask_s1 after conv0, mask_s2 after conv2). +// T_max along time. The convs carry bias + ReLU, so a padded zero becomes +// non-zero after the first conv and then leaks into each utterance's last VALID +// frame via the next stride-2 conv's receptive field. NeMo's masked-subsampling +// fix is to zero each conv intermediate's padded time region after every ReLU +// (mask_s1 after conv0, mask_s2 after conv2). ggml_tensor * build_pre_encode(ggml_context * ctx, const GigaamPreEncode & pe, ggml_tensor * mel_in, @@ -174,11 +147,8 @@ ggml_tensor * build_rotary_attn(ggml_context * ctx, // ne[2] as the position axis (== T) and ne[3] as the batch (== B). ggml_tensor * x_rot = ggml_reshape_4d(ctx, x, head_dim, n_head, T, Bb); // GigaAM passes pos_emb_max_len (default 5000) as the rotary `base` - // (theta). That is NOT the standard 10000 — see - // RotaryPositionalEmbedding(d_model // n_heads, pos_emb_max_len) and - // PositionalEncoding.__init__(dim, base). Verified empirically: the - // dumped enc.pos_emb tensor matches base=5000 to ~1e-6 and base=10000 - // is off by ~0.1. + // (theta), NOT the standard 10000: the dumped enc.pos_emb tensor matches + // base=5000 to ~1e-6 and base=10000 is off by ~0.1. x_rot = ggml_rope_ext(ctx, x_rot, positions, /*c=*/nullptr, /*n_dims=*/head_dim, GGML_ROPE_TYPE_NEOX, @@ -400,7 +370,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_enc); ggml_set_name(positions, "enc.positions"); ggml_set_input(positions); - eb.dumps.pos_emb = positions; // documentation only + eb.dumps.pos_emb = positions; // Conv policy. GigaAM's depthwise k=5 — Metal direct path works. conf::ConvPolicy policy {}; @@ -415,7 +385,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const int n_layers = hp.enc_n_layers; eb.dumps.all_block_outs.reserve(n_layers); - eb.dumps.mid_block_idx = n_layers / 2 - 1; // matches dumper (8 for 16-layer) + eb.dumps.mid_block_idx = n_layers / 2 - 1; eb.dumps.last_block_idx = n_layers - 1; for (int i = 0; i < n_layers; ++i) { @@ -461,8 +431,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // `audio_signal.transpose(1, 2)` — [B, T, D] -> [B, D, T]. Apply // the final transpose to match the reference enc.out byte layout. // CRITICAL: do NOT wrap in a reshape view afterwards; ggml_set_output - // on a view doesn't preserve the materialized cont buffer (verified - // experimentally during M2 bring-up of pre_encode). + // on a view doesn't preserve the materialized cont buffer. ggml_tensor * enc_out_t = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); ggml_set_name(enc_out_t, "enc.out"); eb.out = enc_out_t; diff --git a/src/arch/gigaam/encoder.h b/src/arch/gigaam/encoder.h index 525e127d..06c3e6a2 100644 --- a/src/arch/gigaam/encoder.h +++ b/src/arch/gigaam/encoder.h @@ -1,8 +1,4 @@ // arch/gigaam/encoder.h - GigaAM Conformer encoder graph builder. -// -// M1 stub: build_encoder_graph is not yet implemented. M2 lands the -// pre-encode + rotary PE bank + all 16 conformer blocks. M3 wires the -// per-call run() against this. #pragma once diff --git a/src/arch/gigaam/gigaam.h b/src/arch/gigaam/gigaam.h index 476cc469..15793d32 100644 --- a/src/arch/gigaam/gigaam.h +++ b/src/arch/gigaam/gigaam.h @@ -1,17 +1,11 @@ // arch/gigaam/gigaam.h - GigaAM-family internal model and context types. // -// INTERNAL to src/arch/gigaam/. The public C ABI in include/transcribe.h -// knows nothing about Arch; dispatch happens through find_arch() in -// transcribe-arch.cpp. -// -// GigaAM-v3 ships four ported variants (e2e_rnnt, e2e_ctc, rnnt, ctc). -// All four share a 16-layer Conformer encoder with rotary attention. -// The e2e variants use a SentencePiece tokenizer with punctuation + -// Cyrillic casing; the non-e2e (rnnt, ctc) variants use a charwise -// 33-entry vocab. The upstream `v3_ssl` (HuBERT-CTC pretraining -// checkpoint) is intentionally out of scope: encoder-only, no head, no -// transcribe path, and transcribe.cpp has no encoder-output emission -// CLI. +// INTERNAL to src/arch/gigaam/. GigaAM-v3 ships four ported variants +// (e2e_rnnt, e2e_ctc, rnnt, ctc), all sharing a 16-layer Conformer encoder +// with rotary attention. The e2e variants use a SentencePiece tokenizer +// (punctuation + Cyrillic casing); the non-e2e variants use a charwise +// 33-entry vocab. The upstream `v3_ssl` SSL checkpoint is out of scope +// (encoder-only, no head). #pragma once @@ -44,9 +38,7 @@ namespace transcribe::gigaam { void apply_family_invariants(transcribe_model & model); // Concrete model. Owns the ggml_context that holds every weight tensor's -// data buffer, and any per-family precomputed buffers we may add later -// (Stage 4 milestones will introduce a host decoder mirror analogous to -// parakeet's `host_decoder` for the RNN-T greedy decode loop). +// data buffer plus the host decoder mirror for the RNN-T/CTC greedy loop. struct GigaamModel final : public transcribe_model { Tokenizer tok; GigaamHParams hparams; @@ -61,7 +53,7 @@ struct GigaamModel final : public transcribe_model { GigaamMelFrontend mel; // Host mirror of predictor + joint (RNN-T variants) or CTC head - // (CTC variants). Populated in M3/M4. Empty for now. + // (CTC variants). Populated at load() time. HostDecoderWeights host_decoder; GigaamModel() = default; @@ -81,9 +73,6 @@ struct GigaamSession final : public transcribe_session { std::vector mel_buf; std::vector pos_buf; // cos/sin rotary PE bank, host scratch std::vector enc_host; - // RNN-T greedy decoder scratch goes here when M3 lands; CTC greedy - // scratch lands at M4. - GigaamSession() = default; ~GigaamSession() override; diff --git a/src/arch/gigaam/mel.cpp b/src/arch/gigaam/mel.cpp index 78c50e3a..d70a0670 100644 --- a/src/arch/gigaam/mel.cpp +++ b/src/arch/gigaam/mel.cpp @@ -1,13 +1,9 @@ // arch/gigaam/mel.cpp - GigaAM log-mel feature extractor. // -// Reference: gigaam.preprocess.FeatureExtractor (torchaudio -// MelSpectrogram with mel_scale=htk, n_mels=64, n_fft=win_length=320, -// hop=160, center=False, power=2, periodic Hann), followed by -// SpecScaler = log(clamp(x, 1e-9, 1e9)). -// -// Self-contained: a 320-point FFT for n_fft=320 = 2^6 * 5 uses a -// mixed-radix Cooley-Tukey (radix-2 down to N=5, naive DFT leaf at -// N=5). The same shape transcribe-mel.cpp uses for whisper's n_fft=400. +// Reference: gigaam.preprocess.FeatureExtractor (torchaudio MelSpectrogram, +// htk, n_mels=64, n_fft=win_length=320, hop=160, center=False, power=2, +// periodic Hann) followed by log(clamp(x, 1e-9, 1e9)). n_fft=320 = 2^6 * 5 +// uses a mixed-radix Cooley-Tukey FFT (radix-2 down to a naive DFT leaf at N=5). #include "mel.h" diff --git a/src/arch/gigaam/mel.h b/src/arch/gigaam/mel.h index 03ce4a76..40d3d457 100644 --- a/src/arch/gigaam/mel.h +++ b/src/arch/gigaam/mel.h @@ -17,15 +17,10 @@ namespace transcribe::gigaam { struct GigaamHParams; struct GigaamWeights; -// Load the baked HTK mel filterbank + periodic Hann window from the -// GGUF at load time. hp drives the scalar shape parameters. -// -// The filterbank and window are emitted by the converter from the -// upstream torchaudio state_dict (`preprocessor.featurizer.0.*`), so -// the C++ MelFrontend uses bit-identical buffers — recomputing them -// from htk_hz_to_mel / Hann formulas drifts at high-freq bins because -// torchaudio uses Float32 throughout while a C++ recomputation in -// fp64 produces slightly different bin boundaries. +// Loads the HTK mel filterbank + periodic Hann window baked into the GGUF at +// load time (hp drives the scalar shapes). We use the baked buffers rather than +// recomputing from htk_hz_to_mel / Hann: torchaudio uses Float32 throughout, so +// an fp64 C++ recomputation drifts at high-freq bin boundaries. struct GigaamMelFrontend { int n_freq; // n_fft/2 + 1 int n_mels; diff --git a/src/arch/gigaam/model.cpp b/src/arch/gigaam/model.cpp index 8c2092f2..9e0b2bd6 100644 --- a/src/arch/gigaam/model.cpp +++ b/src/arch/gigaam/model.cpp @@ -1,17 +1,8 @@ // arch/gigaam/model.cpp - GigaAM family handler. // -// M1 (this turn): load() reads hparams, validates the tensor catalog, -// streams weight bytes into a backend buffer, ingests the tokenizer -// and languages KV. init_context() works. run() returns NOT_IMPLEMENTED -// gracefully — the encoder/decoder forward lands in M2/M3. -// -// Stage 4 follow-ups: -// M2: encoder.cpp build_encoder_graph + run() wires it up. -// M3: per-family mel frontend (center=false, htk, log-clamp) and -// RNN-T predictor + joint + greedy decode on host. host_decoder -// populated here at load() time. -// M4: CTC head + greedy collapse. Charwise tokenizer extension lands -// on the shared Tokenizer (not in this file). +// load() reads hparams, validates the tensor catalog, streams weights into a +// backend buffer, and builds the mel frontend + host decoder mirror. run() +// drives the mel -> encoder graph -> host RNN-T/CTC greedy decode pipeline. #include "gigaam.h" @@ -87,23 +78,16 @@ namespace { constexpr const char k_default_variant[] = "gigaam-v3-e2e-rnnt"; -// --------------------------------------------------------------------------- // Input-length contract (see docs/input-limits.md). GigaAM is a SOFT-WINDOW -// family: the Conformer encoder has no hard architectural cap (the rotary -// positional table is recomputed per run), but every published variant was -// trained on utterances up to ~25 s. Beyond that, accuracy degrades rather -// than failing. Upstream gigaam rejects over-length audio outright ("Too long -// wav file"); transcribe.cpp deliberately WARNS-and-PROCEEDS instead so the -// caller keeps control of the degradation. No input gate, no decode/encode -// numeric change — we only warn once past the window. -// --------------------------------------------------------------------------- - -// Advisory training window. No hparam encodes it (the GGUF has no -// max-utterance KV), so we hardcode the upstream "Too long wav file" limit of -// 25 s. k_safe_audio_ms is reported via transcribe_capabilities::max_audio_ms -// and is the threshold for the run() soft-window WARN; the run() check derives -// the clip's duration in ms from n_samples and the frontend sample rate. -// k_safe_s is the same value in seconds, used only in the WARN text. +// family: the Conformer encoder has no hard architectural cap, but every +// published variant was trained on utterances up to ~25 s, beyond which +// accuracy degrades. Upstream rejects over-length audio outright; we +// WARN-and-PROCEED so the caller keeps control of the degradation. +// +// No hparam encodes the window, so hardcode the upstream 25 s limit. +// k_safe_audio_ms is reported via transcribe_capabilities::max_audio_ms and is +// the run() soft-window WARN threshold; k_safe_s is the same value in seconds, +// used only in the WARN text. constexpr int k_safe_s = 25; constexpr int64_t k_safe_audio_ms = 25000; @@ -139,7 +123,7 @@ transcribe_status load(Loader & loader, // than rejecting; see docs/input-limits.md. m->caps.max_audio_ms = k_safe_audio_ms; - // Stage 2: reopen for tensor metadata. + // Reopen for tensor metadata. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -187,10 +171,6 @@ transcribe_status load(Loader & loader, } gguf_free(gguf_data); - // MelFrontend construction: deferred to M3. GigaAM's center=false + - // htk + log-clamp combination doesn't fit the existing MelConfig - // shape; the per-family path lands with the encoder bring-up. - // Build host mirror of the predictor + joint (RNN-T) or CTC head (CTC). if (auto st = build_host_decoder_weights(m->weights, m->hparams, m->host_decoder); @@ -316,7 +296,6 @@ transcribe_status decode_and_populate(GigaamSession * gc, } gc->t_decode_us = ggml_time_us() - t_dec_start; - // Detokenize via the family-agnostic Tokenizer. std::string text = gm->tok.decode(tokens.data(), static_cast(tokens.size())); // SentencePiece convention: the first token's leading ▁ is stripped on @@ -366,11 +345,9 @@ transcribe_status run(transcribe_session * session, transcribe::debug::init(); - // ----- Soft-window advisory (see docs/input-limits.md) --------------- - // GigaAM has no hard cap, but accuracy degrades past the ~25 s window it - // was trained on. Warn once and proceed unchanged — never reject, never - // alter the numerics. The duration is known here from n_samples and the - // frontend sample rate, before any heavy compute. + // Soft-window advisory: warn once past the ~25 s training window and + // proceed unchanged (never reject, never alter numerics). See + // docs/input-limits.md. { const int sr = gm->hparams.fe_sample_rate; if (sr > 0) { @@ -389,10 +366,9 @@ transcribe_status run(transcribe_session * session, } } - // ----- Mel acquisition ----------------------------------------------- - // Production path: the family mel (HTK + power=2 + log-clamp + - // center=False + periodic Hann). The env-var injection stays as a - // debug knob for tensor-parity isolation. + // Mel acquisition: the family mel (HTK + power=2 + log-clamp + + // center=False + periodic Hann). The env-var injection is a debug knob + // for tensor-parity isolation. int mel_n_frames = 0; const int64_t t_mel_start = ggml_time_us(); if (const char * ref_dir = std::getenv("TRANSCRIBE_GIGAAM_MEL_FROM_REF"); @@ -419,7 +395,7 @@ transcribe_status run(transcribe_session * session, if (mel_n_frames <= 0) return TRANSCRIBE_ERR_GGUF; - // ----- Reset per-call compute state --------------------------------- + // Reset per-call compute state. if (gc->compute_ctx != nullptr) { ggml_free(gc->compute_ctx); gc->compute_ctx = nullptr; @@ -441,7 +417,7 @@ transcribe_status run(transcribe_session * session, } } - // ----- Build encoder graph ------------------------------------------ + // Build encoder graph. EncoderBuild eb = build_encoder_graph(gc->compute_ctx, gm->weights, gm->hparams, mel_n_frames, @@ -451,7 +427,7 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_GGUF; } - // ----- Scheduler ---------------------------------------------------- + // Scheduler. if (gc->sched == nullptr) { gc->sched = ggml_backend_sched_new( gm->plan.scheduler_list.data(), nullptr, @@ -474,7 +450,6 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_ERR_OOM; } - // Upload mel. ggml_backend_tensor_set(eb.mel_in, gc->mel_buf.data(), 0, gc->mel_buf.size() * sizeof(float)); transcribe::debug::dump_tensor("enc.mel.in", eb.mel_in, "encoder.mel"); @@ -488,7 +463,7 @@ transcribe_status run(transcribe_session * session, pos.size() * sizeof(int32_t)); } - // ----- Compute ------------------------------------------------------- + // Compute. const int64_t t_enc_start = ggml_time_us(); if (ggml_backend_sched_graph_compute(gc->sched, eb.graph) != GGML_STATUS_SUCCESS) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "gigaam: graph_compute failed"); @@ -496,7 +471,7 @@ transcribe_status run(transcribe_session * session, } gc->t_encode_us = ggml_time_us() - t_enc_start; - // ----- Dump intermediates ------------------------------------------- + // Dump intermediates. if (eb.dumps.pre_encode_out) transcribe::debug::dump_tensor("enc.subsample.out", eb.dumps.pre_encode_out, "encoder.subsample"); @@ -527,13 +502,10 @@ transcribe_status run(transcribe_session * session, gc->encoder_out = eb.out; - // ----- Decoder (RNN-T or CTC greedy) ------------------------------- - // Both heads consume the pre-final-transpose encoder tensor - // (ne=[d_model, T_enc, 1]), which is named `rnnt.encoded` on the - // graph for historical reasons. For CTC variants we read it back - // but skip the `rnnt.encoded` debug dump (the reference dumper does - // not emit it for CTC variants, so emitting one on the C++ side - // would surface as a MISSING-right gate failure in compare_tensors). + // Decoder (RNN-T or CTC greedy). Both heads consume the + // pre-final-transpose encoder tensor (ne=[d_model, T_enc, 1]), named + // `rnnt.encoded` on the graph. CTC variants skip the `rnnt.encoded` dump + // (the reference dumper does not emit it for CTC). { ggml_tensor * enc_t = eb.dumps.rnnt_encoded; if (enc_t == nullptr) { @@ -546,8 +518,6 @@ transcribe_status run(transcribe_session * session, ggml_backend_tensor_get(enc_t, gc->enc_host.data(), 0, gc->enc_host.size() * sizeof(float)); - // Reference `rnnt.encoded` dump (RNNT variants only; the reference - // dumper does not emit it for CTC). validate.py gates on this. if (gm->hparams.head_kind == HeadKind::RNNT) { const long long enc_shape[2] = {T_enc, D}; transcribe::debug::dump_host_f32( @@ -564,20 +534,12 @@ transcribe_status run(transcribe_session * session, return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Offline batch (transcribe_run_batch) -// --------------------------------------------------------------------------- -// -// Batches B utterances through ONE Conformer encoder dispatch (the batch -// rides the activation's ne[2] axis — see encoder.cpp), then host-decodes -// each utterance's RNN-T / CTC slice. The mel front-end and decode are -// identical to single-shot; only the encoder is fused. GigaAM's RNN-T/CTC -// decode is host-side and cheap (small vocab, no BPE-head bottleneck), so — -// unlike granite_nar — encoder batching is the win. Same-length batches run -// mask-free (bit-identical to single-shot per utterance); variable-length -// batches pad to T_max and apply per-utterance masks. Mirrors parakeet's -// run_batch recipe (full-read encoder output + host-slice — NOT per-utt -// offset reads, which are unreliable across backends). +// Offline batch (transcribe_run_batch): B utterances through ONE Conformer +// encoder dispatch (batch on the activation's ne[2] axis), then host-decode +// each utterance's RNN-T / CTC slice. Same-length batches run mask-free +// (bit-identical to single-shot per utterance); variable-length batches pad +// to T_max and apply per-utterance masks. Full-read encoder output + +// host-slice (NOT per-utt offset reads, which are unreliable across backends). // One stride-2 conv with symmetric (k-1)/2 padding (matches encoder.cpp). static int batch_pre_encode_t_out(int in) { return (in - 1) / 2 + 1; } @@ -710,8 +672,7 @@ transcribe_status run_batch_encode(GigaamSession * gc, } gc->t_encode_us = ggml_time_us() - t_enc_start; - // Bisect dump (debug only): full [d, T, n] intermediates so a batched-vs- - // single divergence can be localized per stage. Gated on the env var. + // Debug-only full [d, T, n] intermediates, gated on the env var. if (transcribe::debug::enabled() && std::getenv("TRANSCRIBE_DUMP_ALL_BLOCKS") != nullptr) { @@ -736,7 +697,7 @@ transcribe_status run_batch_encode(GigaamSession * gc, } // Full-read the encoder output, then host-slice per utterance (non-zero - // offset reads are unreliable across backends — see batching-plan.md). + // offset reads are unreliable across backends). const size_t utt_elems = static_cast(d_enc) * static_cast(T_enc); gc->enc_host.assign(utt_elems * static_cast(n), 0.0f); ggml_backend_tensor_get(enc_t, gc->enc_host.data(), 0, @@ -790,10 +751,8 @@ transcribe_status run_batch(transcribe_session * session, const int64_t total_mel_us = ggml_time_us() - t_mel_start; if (all_ok) { - // ----- Soft-window advisory (see docs/input-limits.md) ----------- - // Same warn-and-proceed contract as single-shot run(), applied per - // utterance before the shared fast-path encode. The per-utterance - // fallback below re-enters run(), which warns there, so each + // Per-utterance soft-window advisory before the shared encode; the + // fallback path re-enters run(), which warns there, so each // over-length clip is warned about exactly once. const int sr = gm->hparams.fe_sample_rate; if (sr > 0) { diff --git a/src/arch/gigaam/weights.cpp b/src/arch/gigaam/weights.cpp index dc057741..a9633991 100644 --- a/src/arch/gigaam/weights.cpp +++ b/src/arch/gigaam/weights.cpp @@ -190,7 +190,6 @@ transcribe_status build_block(ggml_context * ctx_meta, const int dff = hp.enc_d_ff; const int dw_k = hp.enc_conv_kernel; - // Macaron FF1. GET_LN (b.norm_ff1_w, lname("enc.blocks.%d.norm_ff1.weight", i)); GET_LN (b.norm_ff1_b, lname("enc.blocks.%d.norm_ff1.bias", i)); GET_LIN_W(b.ff1_lin1_w, lname("enc.blocks.%d.ff1.linear1.weight", i), d, dff); @@ -198,7 +197,6 @@ transcribe_status build_block(ggml_context * ctx_meta, GET_LIN_W(b.ff1_lin2_w, lname("enc.blocks.%d.ff1.linear2.weight", i), dff, d); GET_LIN_B(b.ff1_lin2_b, lname("enc.blocks.%d.ff1.linear2.bias", i), d); - // Conv module. GET_LN (b.norm_conv_w, lname("enc.blocks.%d.norm_conv.weight", i)); GET_LN (b.norm_conv_b, lname("enc.blocks.%d.norm_conv.bias", i)); GET_CONV_W(b.conv_pw1_w, lname("enc.blocks.%d.conv.pointwise1.weight", i), 1, d, 2 * d); @@ -222,7 +220,6 @@ transcribe_status build_block(ggml_context * ctx_meta, GET_LIN_W(b.attn_out_w, lname("enc.blocks.%d.attn.linear_out.weight", i), d, d); GET_LIN_B(b.attn_out_b, lname("enc.blocks.%d.attn.linear_out.bias", i), d); - // Macaron FF2. GET_LN (b.norm_ff2_w, lname("enc.blocks.%d.norm_ff2.weight", i)); GET_LN (b.norm_ff2_b, lname("enc.blocks.%d.norm_ff2.bias", i)); GET_LIN_W(b.ff2_lin1_w, lname("enc.blocks.%d.ff2.linear1.weight", i), d, dff); diff --git a/src/arch/gigaam/weights.h b/src/arch/gigaam/weights.h index b1f6eccd..cc16f272 100644 --- a/src/arch/gigaam/weights.h +++ b/src/arch/gigaam/weights.h @@ -18,7 +18,7 @@ // - Frontend: htk mel scale, no slaney norm. `frontend.mel_filterbank` // (shape [n_mels=64, n_freq_bins=161]) and `frontend.window` // (length 320, Hann periodic) are baked into the GGUF — the C++ -// loader holds borrowed pointers and the M3 mel path reads them +// loader holds borrowed pointers and the mel path reads them // directly. #pragma once diff --git a/src/arch/granite/decoder.cpp b/src/arch/granite/decoder.cpp index aae771df..5e190ab7 100644 --- a/src/arch/granite/decoder.cpp +++ b/src/arch/granite/decoder.cpp @@ -384,7 +384,6 @@ ggml_tensor * block_prefill_batched( Q = ggml_reshape_4d(ctx, Q, head_dim, n_heads, T, B); K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T, B); V = ggml_reshape_4d(ctx, V, head_dim, n_kv_heads, T, B); - // No per-head Q/K RMSNorm on Granite. Q = ggml_rope_ext(ctx, Q, positions, nullptr, static_cast(head_dim), GGML_ROPE_TYPE_NEOX, params.max_position, rope_theta, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); @@ -523,9 +522,7 @@ ggml_tensor * block_step_batched( } // namespace -// --------------------------------------------------------------------------- -// Prefill graph -// --------------------------------------------------------------------------- +// Prefill graph. PrefillBuild build_prefill_graph(ggml_context * ctx, const GraniteWeights & weights, @@ -579,7 +576,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, const auto params = to_params(hp); - // ---------- Graph inputs ---------- + // Graph inputs. pb.input_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_prompt); named(pb.input_ids_in, "dec.input_ids"); ggml_set_input(pb.input_ids_in); @@ -605,14 +602,14 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, } pb.graph = gf; - // ---------- Token embedding (for full prompt, for dump parity) ---------- + // Token embedding (for full prompt, for dump parity). ggml_tensor * token_emb_all = ggml_get_rows( ctx, weights.dec_embed.token_w, pb.input_ids_in); named(token_emb_all, "dec.token_emb"); pb.dumps.token_emb = token_emb_all; transcribe::debug::mark_tensor_for_dump(token_emb_all); - // ---------- Audio injection via 3-way concat ---------- + // Audio injection via 3-way concat. const size_t emb_elem = ggml_element_size(token_emb_all); ggml_tensor * x_prefix = nullptr; if (prefix_len > 0) { @@ -642,10 +639,10 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, pb.dumps.audio_injected = x; transcribe::debug::mark_tensor_for_dump(x); - // ---------- Embedding multiplier (after the dump point) ---------- + // Embedding multiplier (after the dump point). x = ggml_scale(ctx, x, emb_mul); - // ---------- Block stack ---------- + // Block stack. const int mid_idx = n_layer / 2; for (int il = 0; il < n_layer; ++il) { x = block_prefill( @@ -678,13 +675,13 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, } } - // ---------- Final RMSNorm ---------- + // Final RMSNorm. x = rms_norm(ctx, x, weights.dec_final.norm_w, rms_eps); named(x, "dec.out_before_head"); pb.dumps.out_before_head = x; transcribe::debug::mark_tensor_for_dump(x); - // ---------- LM head (slice last position) ---------- + // LM head (slice last position). ggml_tensor * last_x = ggml_view_2d( ctx, x, hidden, 1, @@ -720,9 +717,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, return pb; } -// --------------------------------------------------------------------------- -// Step graph -// --------------------------------------------------------------------------- +// Step graph. StepBuild build_step_graph(ggml_context * ctx, const GraniteWeights & weights, @@ -823,9 +818,7 @@ StepBuild build_step_graph(ggml_context * ctx, return sb; } -// --------------------------------------------------------------------------- -// Batched prefill / step (offline transcribe_run_batch) -// --------------------------------------------------------------------------- +// Batched prefill / step (offline transcribe_run_batch). PrefillBuildBatched build_prefill_graph_batched( ggml_context * ctx, diff --git a/src/arch/granite/decoder.h b/src/arch/granite/decoder.h index 263fcd5d..18e2c033 100644 --- a/src/arch/granite/decoder.h +++ b/src/arch/granite/decoder.h @@ -11,7 +11,14 @@ // - Attention softmax scale = attention_multiplier (1/128), NOT // 1/sqrt(head_dim). // - Residual adds use residual_multiplier (0.22): x + 0.22 * sub_out -// - SwiGLU MLP (separate gate + up; no packing in first port) +// - SwiGLU MLP. When ffn_gate_up_w is present (packed at load via +// causal_lm::pack_gate_up), the step, batched-prefill, and +// batched-step paths take the fused single-mul_mat + ggml_swiglu +// route, falling back to separate gate/up mul_mats otherwise. The +// single (non-batched) prefill path always uses separate gate + up. +// See decoder.cpp (ffn_gate_up_w branches in block_step / +// block_prefill_batched / block_step_batched; block_prefill is +// unconditionally separate). // // Outer flow (matches HF GraniteSpeechForConditionalGeneration + // GraniteModel.forward): @@ -22,7 +29,7 @@ // 5. x = RMSNorm(x) # dec.out_before_head // 6. logits = lm_head(x) / logits_scaling # dec.logits_raw (last pos) // -// First port handles a single contiguous audio block at positions +// The public path handles a single contiguous audio block at positions // [prefix_len, prefix_len + n_audio_tokens). Granite's chat template // has exactly one <|audio|> placeholder per turn so multi-audio prompts // aren't reachable from the public API today. @@ -85,7 +92,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, bool use_flash, bool slice_last); -// ---------- Step graph (one token) ---------- +// Step graph (one token). struct StepBuild { ggml_tensor * input_id_in = nullptr; // [1] i32 @@ -107,7 +114,7 @@ StepBuild build_step_graph(ggml_context * ctx, int max_n_kv, bool use_flash); -// ---------- Batched prefill / step (offline transcribe_run_batch) ---------- +// Batched prefill / step (offline transcribe_run_batch). // Same recipe as the causal_lm families but with Granite block math (no Q/K // norm, attention_multiplier scale, residual_multiplier, embedding_multiplier, // logits_scaling). Batch rides ne[2]; requires use_flash. diff --git a/src/arch/granite/encoder.cpp b/src/arch/granite/encoder.cpp index 6f47a5d2..14971993 100644 --- a/src/arch/granite/encoder.cpp +++ b/src/arch/granite/encoder.cpp @@ -48,9 +48,7 @@ namespace transcribe::granite { -// --------------------------------------------------------------------------- -// Constants -// --------------------------------------------------------------------------- +// Constants. namespace { @@ -88,9 +86,7 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, } // namespace -// --------------------------------------------------------------------------- -// Host-side mel + 2-frame stack -// --------------------------------------------------------------------------- +// Host-side mel + 2-frame stack. transcribe_status compute_mel_encoder_input( const transcribe::MelFrontend & mel, @@ -154,9 +150,7 @@ transcribe_status compute_mel_encoder_input( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Shaw attention_dists + last-block mask -// --------------------------------------------------------------------------- +// Shaw attention_dists + last-block mask. std::vector precompute_attention_dists(int context_size, int max_pos_emb) { @@ -167,12 +161,9 @@ std::vector precompute_attention_dists(int context_size, int max_pos_em // // `seq.view(-1, 1) - seq.view(1, -1)` broadcasts to a [c, r] matrix // whose element at (c, r) is `c - r` — i.e., the offset from the - // KEY position back to the QUERY position. (i in the row index is - // the query position; j in the column index is the key position.) - // - // We were writing `r - c` here — the opposite sign — which mirrored - // the Shaw bias across the diagonal and corrupted block-local - // attention from the first encoder layer. + // KEY position back to the QUERY position. (Row index = query position; + // column index = key position. The sign matters: `r - c` would mirror + // the Shaw bias across the diagonal.) std::vector dists(static_cast(context_size) * context_size); for (int c = 0; c < context_size; ++c) { for (int r = 0; r < context_size; ++r) { @@ -216,9 +207,7 @@ std::vector precompute_last_block_mask(int context_size, int t_enc_remain return mask; } -// --------------------------------------------------------------------------- -// Per-block builders -// --------------------------------------------------------------------------- +// Per-block builders. namespace { @@ -250,14 +239,10 @@ ggml_tensor * granite_macaron(ggml_context * ctx, // Conv module: LN -> pointwise1 (d_model -> 2*inner_dim) -> GLU split // (-> inner_dim) -> depthwise (k=15) -> BN -> SiLU -> pointwise2 -// (inner_dim -> d_model). Operates on [d_model, T] (ne ordering). -// -// inner_dim = d_model * conv_expansion. For granite conv_expansion=2, -// so inner_dim = 2*d_model and the pw1 kernel maps d_model -> 4*d_model. -// -// The conformer::conv_module helper hardcodes inner_dim == d_model -// (parakeet / cohere / canary all use conv_expansion=1). We inline the -// module here. The shape of operations is otherwise identical. +// (inner_dim -> d_model), on [d_model, T]. inner_dim = d_model * +// conv_expansion (=2 on granite). Inlined rather than using +// conformer::conv_module, which hardcodes inner_dim == d_model +// (conv_expansion=1 on parakeet / cohere / canary). ggml_tensor * granite_conv_module(ggml_context * ctx, ggml_tensor * x, const GraniteEncBlock & b, @@ -309,7 +294,6 @@ ggml_tensor * granite_conv_module(ggml_context * ctx, bn_fused_scale, bn_fused_bias); - // SiLU. x = ggml_silu(ctx, x); // Transpose back: [T, inner_dim] -> [inner_dim, T]. @@ -331,9 +315,7 @@ ggml_tensor * granite_conv_module(ggml_context * ctx, } // namespace -// --------------------------------------------------------------------------- -// Encoder graph -// --------------------------------------------------------------------------- +// Encoder graph. EncoderBuild build_encoder_graph(ggml_context * ctx, const GraniteWeights & weights, @@ -355,7 +337,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const int head_dim = hp.enc_head_dim; const int ctx_size = hp.enc_context_size; - // ----- Graph inputs ----- + // Graph inputs. // mel_in: [input_dim, T_enc]. Encoder operates at T_enc throughout; // the only T_pad expansion happens INSIDE the Shaw block-local // attention helper, mirroring the reference's @@ -390,7 +372,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_set_input(eb.zero_pad); } - // ----- Input linear (160 -> 1024) ----- + // Input linear (160 -> 1024). ggml_tensor * x = linear(ctx, eb.mel_in, weights.enc_top.input_linear_w, weights.enc_top.input_linear_b); @@ -399,26 +381,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.input_linear_out = x; transcribe::debug::mark_tensor_for_dump(x); - // BN fusion happens at load time and we read scale/bias from the - // model (filled in model.cpp). For Phase 2 we accept that the - // fused tensors come from external storage (via build_encoder_graph - // caller wiring); we look them up off the granite weights struct. - // - // Each block in granite has its own [inner_dim] BN gamma/beta plus - // running stats. The graph references the FUSED scale/bias which - // model.cpp computes and writes once. They live as separate - // ggml_tensor*s outside the GraniteWeights struct (in - // GraniteModel::bn_fused); see model.cpp::fuse_batch_norm(). - // - // For Phase 2 we accept a callback-style override: the caller must - // ensure GraniteEncBlock::conv_bn_w / conv_bn_b / conv_bn_mean / - // conv_bn_var have been "fused" into a pair of [inner_dim] f32 - // tensors and exposed through a side table. To keep build_encoder - // self-contained we recompute the fusion on the fly here using - // ggml ops: + // The graph references per-block FUSED BN scale/bias that model.cpp + // computes once at load time (see model.cpp::fuse_batch_norm): // scale_i = bn_gamma_i / sqrt(bn_var_i + eps) // bias_i = bn_beta_i - bn_mean_i * scale_i - // These are 1-D fp32 tensors that fold cleanly into the graph. const int n_layers = static_cast(weights.enc_blocks.size()); const int bypass_after = hp.enc_n_layers / 2; // 1-indexed boundary @@ -519,10 +485,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, transcribe::debug::mark_tensor_for_dump(x); } if (i == bypass_after) { - // i.e., block index N/2 (0-indexed). For n=16, this is - // block 8 — Stage 2 dumps `enc.block.8.out` after block 8's - // forward completes (with the bypass residual already - // injected at block-8 input). + // Block index N/2 (0-indexed); for n=16 this is block 8, dumped + // after its forward completes (bypass residual already injected). char bname[64]; std::snprintf(bname, sizeof(bname), "enc.block.%d.out", i); named(x, bname); @@ -538,7 +502,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // ----- cat_hidden_layers channel concat (-plus only) ----- + // cat_hidden_layers channel concat (-plus only). // HF: hidden_states = torch.cat([*exported_hidden_states, hidden_states], dim=-1) // i.e. earlier-layer captures FIRST, final hidden LAST, along the // channel axis. ggml dim=0 is the channel axis for [hidden, T]. @@ -550,7 +514,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, x = ggml_concat(ctx, *it, x, /*dim=*/0); } - // ----- Final output (already at T_enc) ----- + // Final output (already at T_enc). named(x, "enc.out"); eb.out = x; eb.dumps.out_named = x; diff --git a/src/arch/granite/encoder.h b/src/arch/granite/encoder.h index 46781eff..1ef92be4 100644 --- a/src/arch/granite/encoder.h +++ b/src/arch/granite/encoder.h @@ -38,9 +38,7 @@ namespace transcribe { class MelFrontend; } namespace transcribe::granite { -// --------------------------------------------------------------------------- -// Host-side mel preprocessing -// --------------------------------------------------------------------------- +// Host-side mel preprocessing. // Run the reference MelSpectrogram (via transcribe::MelFrontend in // whisper-mode + htk filterbank + reflect-pad + hann_periodic window), @@ -58,9 +56,7 @@ transcribe_status compute_mel_encoder_input( std::vector & out_mel, int & out_t_enc); -// --------------------------------------------------------------------------- -// Graph construction -// --------------------------------------------------------------------------- +// Graph construction. struct EncoderBuild { // Graph inputs (caller uploads at compute time). @@ -80,7 +76,7 @@ struct EncoderBuild { ggml_cgraph * graph = nullptr; - // Dump points (exposed so model.cpp can wire them up for validate.py). + // Dump points (exposed so model.cpp can wire them up). struct Dumps { ggml_tensor * input_linear_out = nullptr; // [hidden, T_enc] ggml_tensor * block_0_out = nullptr; @@ -97,7 +93,7 @@ struct EncoderBuild { // Build the encoder graph. `T_enc` is the number of post-stack frames // (== n_mel_frames / 2 after the whisper-mode trim). `use_flash` is -// reserved for future use — the first port uses manual mul_mat + soft_max +// 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. EncoderBuild build_encoder_graph(ggml_context * ctx, @@ -107,9 +103,11 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, bool use_flash); // Host-side precomputation of the Shaw attention_dists matrix. -// attention_dists[c, r] = clamp(r - c, -context_size, context_size) + max_pos_emb, -// flattened to row-major int32 [context_size * context_size]. Caller -// uploads this once per encode into EncoderBuild::attention_dists. +// attention_dists[c, r] = clamp(c - r, -context_size, context_size) + max_pos_emb, +// where c is the key/column index and r is the query/row index (matches the +// reference `seq.view(-1,1) - seq.view(1,-1)`; see precompute_attention_dists +// in encoder.cpp). Flattened to row-major int32 [context_size * context_size]. +// Caller uploads this once per encode into EncoderBuild::attention_dists. // // The same matrix is reused across every encoder block and every batch // element (and across decode calls if the encoder shape is unchanged). diff --git a/src/arch/granite/granite.h b/src/arch/granite/granite.h index f1979196..c5fc35ce 100644 --- a/src/arch/granite/granite.h +++ b/src/arch/granite/granite.h @@ -5,8 +5,7 @@ // Speech family (audio-LLM: Conformer encoder + BLIP-2 Q-Former // projector + Granite-4 causal LM with audio-token injection). // -// First port targets non-streaming ASR transcription across 3 AR -// variants: +// Supported non-streaming ASR variants: // - granite-4.0-1b-speech // - granite-speech-4.1-2b // - granite-speech-4.1-2b-plus @@ -53,15 +52,10 @@ namespace transcribe::granite { void apply_family_invariants(transcribe_model & model); -// Chat-template token ids resolved through the loaded tokenizer at -// load time. Granite's bare-USER:/ASSISTANT: template (1b/2b) and the -// Granite-4 system-role template (plus) overlap on the role-marker -// tokens; we resolve every piece we might need ahead of time so a -// future vocab drift fails loudly at load instead of mid-decode. -// -// All fields are the *resolved* ids. The reference IDs on the granite -// vocab (100256..100352 are added tokens) are stable across the 3 AR -// variants, but we look them up rather than hardcode. +// Chat-template token ids resolved through the loaded tokenizer at load time +// (bare-USER:/ASSISTANT: for 1b/2b, Granite-4 system-role for -plus). Every +// piece is resolved ahead of time so vocab drift fails loudly at load instead +// of mid-decode. All fields are resolved ids — looked up, not hardcoded. struct ChatTokens { int32_t audio = -1; // <|audio|> int32_t end_of_text = -1; // <|end_of_text|> (also EOS / BOS) @@ -70,9 +64,7 @@ struct ChatTokens { int32_t end_of_role = -1; // <|end_of_role|> (granite-4 chat only) }; -// --------------------------------------------------------------------------- -// Model / Context -// --------------------------------------------------------------------------- +// Model / Context. struct GraniteModel final : public transcribe_model { Tokenizer tok; @@ -98,10 +90,9 @@ struct GraniteModel final : public transcribe_model { std::optional mel; // Jinja chat template from the GGUF KV. Stored verbatim — we don't - // render Jinja at runtime; the family-specific prompt builder - // (Phase 5) emits the equivalent token sequence directly. Empty - // string means "no template" (the runtime will refuse to decode - // without one). + // render Jinja at runtime; the family-specific prompt builder emits the + // equivalent token sequence directly. Empty string means "no template" + // (the runtime will refuse to decode without one). std::string chat_template; // Resolved chat-template token ids (filled at load time). diff --git a/src/arch/granite/model.cpp b/src/arch/granite/model.cpp index d786b612..0c808279 100644 --- a/src/arch/granite/model.cpp +++ b/src/arch/granite/model.cpp @@ -1,27 +1,4 @@ // arch/granite/model.cpp - Granite Speech family handler. -// -// Wiring skeleton (Phase 1 of porting-4-cpp). -// -// load() is real: reads the GGUF KV, resolves the tokenizer, -// configures the mel frontend, builds the tensor -// catalog, allocates the backend buffer, and streams -// tensor data. This lets `transcribe-cli -m ...` -// load every granite variant and print its hparams -// before any encode/decode code lands. -// init_context() returns TRANSCRIBE_ERR_NOT_IMPLEMENTED. -// run() returns TRANSCRIBE_ERR_NOT_IMPLEMENTED. -// -// Phase 2-5 fill in the encoder graph, projector, decoder LM, and -// end-to-end run(). The Phase 1 contract is "loader smoke succeeds"; -// we deliberately error out on context creation so a caller cannot -// accidentally invoke an empty encode pipeline. -// -// The mel frontend currently builds with the existing MelFrontend -// surface. Granite's torchaudio melspec needs htk-mel norm, no log, -// no per-utterance norm, and a 2-frame stack — those knobs are added -// in Phase 2b. For Phase 1 we just record the cfg and accept whatever -// the existing frontend produces; the encoder isn't wired up yet so -// the discrepancy is not user-visible. #include "granite.h" @@ -108,13 +85,11 @@ namespace { constexpr const char k_default_variant[] = "granite-speech"; constexpr float kBnEps = 1e-5f; -// --------------------------------------------------------------------------- // Input-length contract (see docs/input-limits.md). Granite is a // hard-context-cap family: audio tokens + prompt + generation share the LLM // decoder's context window (dec_max_position_embeddings, typically 4096). // Over-length input is rejected up front with TRANSCRIBE_ERR_INPUT_TOO_LONG // rather than silently aliasing RoPE past the trained range. -// --------------------------------------------------------------------------- // Generation budget reserved per run. Also the KV grow-to-fit step budget, // so an accepted clip always has room for up to this many output tokens. @@ -169,11 +144,8 @@ int64_t granite_max_audio_ms(const GraniteHParams & hp) { return frames * hp.fe_hop_length * 1000 / hp.fe_sample_rate; } -// Pre-fuse BatchNorm scale/bias per encoder block. Mirrors parakeet's -// fuse_batch_norm: allocates a separate ggml context + CPU buffer for -// the [inner_dim] fused tensors, copies the raw BN tensors out, -// computes scale = gamma / sqrt(var + eps) and bias = beta - mean * scale, -// and uploads. +// Pre-fuse BatchNorm scale/bias per encoder block (mirrors parakeet's +// fuse_batch_norm): scale = gamma / sqrt(var + eps), bias = beta - mean * scale. transcribe_status fuse_batch_norm(GraniteModel & m) { const size_t n_blocks = m.weights.enc_blocks.size(); if (n_blocks == 0) return TRANSCRIBE_OK; @@ -235,14 +207,10 @@ transcribe_status fuse_batch_norm(GraniteModel & m) { return TRANSCRIBE_OK; } -// Resolve the chat-template pieces we may need across the two -// templates (bare-USER/ASSISTANT for 1b/2b, granite-4 system-role for -// -plus). The audio token must exist on every variant; the -// start_of_role / end_of_role pieces only exist on the -plus vocab, -// so they're allowed to be absent. -// -// Hard-fails for the must-have pieces (audio, end_of_text, pad) so a -// future vocab drift surfaces at load rather than mid-decode. +// Resolve chat-template pieces across both templates (bare-USER/ASSISTANT for +// 1b/2b, granite-4 system-role for -plus). audio exists on every variant; +// start_of_role / end_of_role only on -plus, so they may be absent. Hard-fails +// for the must-haves (audio, end_of_text, pad) so vocab drift surfaces at load. transcribe_status resolve_chat_tokens(const transcribe::Tokenizer & tok, ChatTokens & out) { @@ -312,7 +280,7 @@ transcribe_status load( // Publish the input-length ceiling now that the decoder context window // and frontend rate are known (apply_family_invariants ran before the - // hparams were read, so it could not set this). See docs/input-limits.md. + // hparams were read, so it could not set this). m->caps.max_audio_ms = granite_max_audio_ms(m->hparams); // Basis for the session-level limits query (transcribe_session_get_limits): @@ -360,13 +328,7 @@ transcribe_status load( return TRANSCRIBE_ERR_GGUF; } - // Mel frontend. Phase 1: we configure with the granite KV but the - // existing MelFrontend doesn't yet honour htk-norm / no-log / - // power=2 / per-utterance-stack. Phase 2b extends the frontend - // surface. Since the encoder graph isn't wired up in Phase 1, the - // frontend's actual output is not observed by any code path that - // matters yet — but we still construct it so the loader smoke - // exercises the same code path the eventual run() will use. + // Mel frontend: configured from the granite KV. { transcribe::MelConfig cfg {}; cfg.sample_rate = m->hparams.fe_sample_rate; @@ -383,7 +345,7 @@ transcribe_status load( // Frontend buffers baked by the converter: librosa htk mel // filterbank + periodic Hann window. We pick them up here so - // the eventual frontend bring-up uses bit-identical values. + // the frontend uses bit-identical values. { using R = transcribe::load_common::ReadF32Result; const size_t fb_elems = @@ -409,7 +371,7 @@ transcribe_status load( m->mel.emplace(cfg); } - // Stage 2: reopen with no_alloc to build the tensor catalog. + // Reopen with no_alloc to build the tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -513,8 +475,7 @@ transcribe_status init_context( // Encoder uses manual mul_mat + soft_max (no flash) because the // Shaw bias requires a per-(head, block) additive term that - // flash_attn_ext doesn't broadcast cleanly yet. Decoder flash is - // wired in Phase 4. + // flash_attn_ext doesn't broadcast cleanly yet. cc->encoder_use_flash = false; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, @@ -524,11 +485,6 @@ transcribe_status init_context( return TRANSCRIBE_OK; } -// Phase 2 run(): mel + 2-frame stack -> encoder graph -> dump enc.* -// tensors. The projector + LM are still NOT_IMPLEMENTED, so we set -// has_result/full_text to "" so validate.py picks up a `text:` line -// and the per-tensor compare can run against enc.* dumps. Phase 3-5 -// replaces this body with the full pipeline. // Map a BCP-47 language code or English name to the language name the // granite-speech instruction expects. Returns nullptr for unsupported. // granite-speech base AR variants advertise fr/de/es/pt/ja plus translate-only @@ -672,7 +628,7 @@ transcribe_status run( return TRANSCRIBE_ERR_INVALID_ARG; } - // ----- Mel + 2-frame stack (host side) ----- + // Mel + 2-frame stack (host side). const int64_t t_mel_start = ggml_time_us(); int t_enc = 0; if (const transcribe_status mst = compute_mel_encoder_input( @@ -697,7 +653,7 @@ transcribe_status run( shape, 2, "encoder"); } - // ----- Reset compute state ----- + // Reset compute state. if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; @@ -717,7 +673,7 @@ transcribe_status run( } } - // ----- Build encoder graph ----- + // Build encoder graph. EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash); @@ -729,7 +685,7 @@ transcribe_status run( // the mel buffer required). The Shaw-attention helper pads to // T_pad internally using the eb.zero_pad graph input. - // ----- Scheduler ----- + // Scheduler. if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new( cm->plan.scheduler_list.data(), nullptr, @@ -748,7 +704,7 @@ transcribe_status run( return TRANSCRIBE_ERR_OOM; } - // ----- Upload inputs ----- + // Upload inputs. // mel_in: contiguous [T_enc, input_dim] row-major matches ggml ne // [input_dim, T_enc] (ne[0]=input_dim is innermost). ggml_backend_tensor_set(eb.mel_in, cc->mel_buf.data(), 0, @@ -791,7 +747,7 @@ transcribe_status run( // Thread count. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); - // ----- Compute ----- + // Compute. const int64_t t_enc_start = ggml_time_us(); if (const ggml_status gs = ggml_backend_sched_graph_compute(cc->sched, eb.graph); @@ -804,7 +760,7 @@ transcribe_status run( } cc->t_encode_us = ggml_time_us() - t_enc_start; - // ----- Dump intermediates ----- + // Dump intermediates. auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { if (t != nullptr) transcribe::debug::dump_tensor(name, t, stage); }; @@ -824,11 +780,9 @@ transcribe_status run( } try_dump("enc.out", eb.dumps.out_named, "encoder"); - // ----- Read encoder output to host for the projector pass ----- - // Phase 3 uses a SEPARATE compute graph for the projector so the - // per-stage validation harness can compare proj.* dumps without - // any encoder graph state lingering. Phase 5 may fuse encoder + - // projector into one compute when end-to-end latency matters. + // Read encoder output to host for the projector pass. The projector + // uses a separate compute graph so its dumps compare without encoder + // graph state lingering. const int64_t enc_hidden_runtime = eb.out->ne[0]; const int64_t enc_T_runtime = eb.out->ne[1]; std::vector enc_host(static_cast(enc_hidden_runtime) * @@ -836,14 +790,12 @@ transcribe_status run( ggml_backend_tensor_get(eb.out, enc_host.data(), 0, enc_host.size() * sizeof(float)); - // ----- Cat-hidden-layers handling (-plus variant) ----- - // The -plus checkpoint sets enc.cat_hidden_layers = [3]: the - // encoder graph (see encoder.cpp) channel-concatenates block[K-1].out - // captures with the final hidden along dim=0 BEFORE returning eb.out, - // so the projector input dimension widens from 1024 to 2048 - // automatically. Nothing variant-specific is needed here. + // Cat-hidden-layers handling (-plus variant): the -plus checkpoint sets + // enc.cat_hidden_layers = [3]; the encoder graph channel-concatenates + // block[K-1].out with the final hidden along dim=0 before returning eb.out, + // so the projector input dimension widens from 1024 to 2048 automatically. - // ----- Projector graph ----- + // Projector graph. ggml_context * proj_ctx = nullptr; { ggml_init_params ip {}; @@ -908,7 +860,7 @@ transcribe_status run( ggml_free(proj_ctx); - // ----- Decoder prefill pass ----- + // Decoder prefill pass. // Build the prompt as `prefix_text | <|audio|>×n | suffix_text`. // Two templates ship today: // @@ -932,9 +884,8 @@ transcribe_status run( // (leading space matters: BPE token IDs differ) // PLUS a Granite-specific system prompt (see use_granite4_chat below). // word-timestamps task : "transcribe the speech with timestamps in [SS:MS] format" - // (only -plus actually emits the [SS:N] markers; 1b/2b - // fall back to plain transcript, which Stage-4 Capability - // Validation records as SKIP rather than PASS) + // (only -plus emits the [SS:N] markers; 1b/2b + // fall back to plain transcript) // translate task : "can you translate the speech into ?" std::vector prefix_ids; std::vector suffix_ids; @@ -949,12 +900,10 @@ transcribe_status run( const int suffix_len = static_cast(suffix_ids.size()); const int T_prompt = prefix_len + n_audio_tokens + suffix_len; - // Reference quirk: HF's GraniteSpeechForConditionalGeneration replaces - // audio_token_id with 0 BEFORE the embed_tokens lookup (the - // resulting rows are scattered over with audio_features moments - // later). We mirror that masking so dec.token_emb matches at the - // audio positions — the forward result is identical either way - // because those rows are overwritten by the audio scatter. + // Reference quirk: HF replaces audio_token_id with 0 before the + // embed_tokens lookup (those rows are overwritten by the audio scatter + // moments later). We mirror the masking so dec.token_emb matches at the + // audio positions; the forward result is identical either way. std::vector input_ids; input_ids.reserve(T_prompt); input_ids.insert(input_ids.end(), prefix_ids.begin(), prefix_ids.end()); @@ -963,11 +912,10 @@ transcribe_status run( } input_ids.insert(input_ids.end(), suffix_ids.begin(), suffix_ids.end()); - // ----- Input-length gate (see docs/input-limits.md) ----- - // The decoder context window is the binding limit: audio tokens + - // prompt + generation must fit dec_max_position_embeddings (optionally - // lowered by the caller's n_ctx knob). The audio-token count is fixed - // by the input length, so reject an over-length clip here, before the + // Input-length gate. The decoder context window is the binding limit: + // audio tokens + prompt + generation must fit dec_max_position_embeddings + // (optionally lowered by the caller's n_ctx knob). The audio-token count is + // fixed by the input length, so reject an over-length clip here, before the // autoregressive decode, instead of growing the cache unboundedly and // aliasing RoPE past the trained range. Reserving the full generation // budget means an accepted clip always has room for a real transcript. @@ -1138,7 +1086,7 @@ transcribe_status run( int32_t next_id = argmax_logits(last_logits); - // ----- Greedy step loop ----- + // Greedy step loop. // Build one reusable step graph sized for the whole run. The // n_ctx of the KV cache bounds the max generation length we can // attend over. @@ -1225,7 +1173,7 @@ transcribe_status run( // The decode stopped either at EOS (complete) or at the generation // budget / context ceiling (truncated). Surface the latter via // transcribe_was_truncated() and a WARN rather than handing back a - // silently shortened transcript. See docs/input-limits.md. + // silently shortened transcript. if (next_id != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -1236,7 +1184,7 @@ transcribe_status run( static_cast(gen_ids.size())); } - // ----- Detokenize ----- + // Detokenize. cc->full_text = cm->tok.decode(gen_ids.data(), static_cast(gen_ids.size())); @@ -1257,13 +1205,11 @@ transcribe_status run( : TRANSCRIBE_OK; } -// =========================================================================== -// Offline batched decode (transcribe_run_batch) -// =========================================================================== -// Serial mel + Conformer encoder + Q-Former projector per utterance produce -// each one's audio embedding [hidden, n_audio_tokens]; then prefill + step are -// batched via granite's batched block builders. Same recipe as the causal_lm -// families, with Granite-specific block math. +// Offline batched decode (transcribe_run_batch). Serial mel + Conformer +// encoder + Q-Former projector per utterance produce each one's audio embedding +// [hidden, n_audio_tokens]; then prefill + step are batched via granite's +// batched block builders. Same recipe as the causal_lm families, with +// Granite-specific block math. transcribe_status reset_ctx_g(GraniteSession * cc, int mb) { if (cc->compute_ctx != nullptr) { @@ -1282,17 +1228,12 @@ void apply_threads_g(GraniteSession * cc) { transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } -// encoder + projector for one utterance from a PRECOMPUTED mel buffer (the -// mel+2-frame-stack is computed in parallel by the caller) → [hidden, n_audio]. -// -// NOTE: the encoder is run PER-UTTERANCE (serial) deliberately. The 2B -// conformer is compute-bound — the time axis (T_enc) already amortizes the -// weight reads and the matmuls saturate the GPU — so a batched single-graph -// encoder is a measured wash on BOTH Metal and L40S (wall time identical -// batched vs serial; encoder compute flat across batch sizes), and for -// mixed-length batches it is strictly worse (it pads every utterance to the -// longest). See docs/batching-autoregressive-plan.md. Only the latency-bound -// decode (M=1) and the host-side mel benefit from batching/threading. +// encoder + projector for one utterance from a PRECOMPUTED mel buffer → +// [hidden, n_audio]. The encoder runs PER-UTTERANCE (serial) deliberately: the +// 2B conformer is compute-bound, so a batched single-graph encoder is a +// measured wash on Metal and L40S and strictly worse for mixed-length batches +// (it pads to the longest). Only the decode and host-side mel benefit from +// batching/threading. transcribe_status encode_one( GraniteSession * cc, GraniteModel * cm, const std::vector & mel_buf, int t_enc, @@ -1451,7 +1392,7 @@ transcribe_status run_batch( return TRANSCRIBE_ERR_INVALID_ARG; const int prefix_len = static_cast(prefix_ids.size()); - // ---- Pass 1: per-utterance mel + encoder + projector (serial) ---- + // Pass 1: per-utterance mel + encoder + projector (serial). std::vector valid(n, 0); // Per-utterance failure status for the result capture below. Defaults to // INVALID_ARG (bad pcm / mel / encode); the input-length gate upgrades it @@ -1463,13 +1404,13 @@ transcribe_status run_batch( std::vector T_prompt(n, 0); int64_t mel_us = 0, enc_us = 0; - // Decoder context ceiling for the per-utterance input-length gate (see - // docs/input-limits.md). Same value the single-utterance run() enforces; - // an over-length utterance is rejected with TRANSCRIBE_ERR_INPUT_TOO_LONG - // rather than growing the cache unboundedly. + // Decoder context ceiling for the per-utterance input-length gate. Same + // value the single-utterance run() enforces; an over-length utterance is + // rejected with TRANSCRIBE_ERR_INPUT_TOO_LONG rather than growing the cache + // unboundedly. const int ceiling = granite_context_ceiling(cc->n_ctx, cm->hparams); - // ---- Pass 0: parallel mel + 2-frame stack (host-side, thread-safe) ---- + // Pass 0: parallel mel + 2-frame stack (host-side, thread-safe). std::vector> mel_bufs(n); std::vector mel_t_enc(n, 0); int n_mel_threads = cc->n_threads; @@ -1486,12 +1427,8 @@ transcribe_status run_batch( }); mel_us += ggml_time_us() - t_mel0; - // ---- Pass 1: per-utterance encoder + projector (serial) ---- - // The conformer encoder is compute-bound (see encode_one's note): a - // batched single-graph encoder is a wash on Metal AND L40S (wall - // identical) and regresses on mixed-length batches, so it is NOT - // batched. Only the mel (Pass 0) is parallelized; the decode (prefill + - // step loop) is batched below. + // Pass 1: per-utterance encoder + projector (serial — see encode_one's + // note). Only the mel (Pass 0) is parallelized; the decode is batched below. for (int b = 0; b < n; ++b) { if (cc->poll_abort()) return TRANSCRIBE_ERR_ABORTED; if (mel_t_enc[b] <= 0) continue; @@ -1504,10 +1441,7 @@ transcribe_status run_batch( prompt_ids[b].insert(prompt_ids[b].end(), suffix_ids.begin(), suffix_ids.end()); T_prompt[b] = static_cast(prompt_ids[b].size()); - // Input-length gate (see docs/input-limits.md). Audio tokens + prompt + - // generation must fit the decoder context window; reject an over-length - // utterance here instead of growing the cache unboundedly. Mirrors the - // single-shot run() gate (T_prompt + k_gen_budget > ceiling). + // Input-length gate, mirroring the single-shot run() gate. if (T_prompt[b] + k_gen_budget > ceiling) { transcribe::log_msg( TRANSCRIBE_LOG_LEVEL_ERROR, @@ -1543,7 +1477,7 @@ transcribe_status run_batch( // ceiling (the per-utterance gate guarantees every valid row fits). if (max_n_kv > ceiling) max_n_kv = ceiling; - // ---- Allocate batched KV cache ---- + // Allocate batched KV cache. ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; if (cc->kv_batch.self_k == nullptr || @@ -1562,7 +1496,7 @@ transcribe_status run_batch( ggml_backend_buffer_clear(cc->kv_batch.buffer, 0); } - // ---- Pass 2: batched prefill ---- + // Pass 2: batched prefill. std::vector next_tok(n, 0); std::vector n_past(n, 0); std::vector> generated(n); @@ -1648,7 +1582,7 @@ transcribe_status run_batch( const int64_t prefill_us = ggml_time_us() - t_pref0; - // ---- Pass 3: batched step loop (shared causal_lm driver) ---- + // Pass 3: batched step loop (shared causal_lm driver). const int32_t eos_id = cm->hparams.eos_token_id; if (reset_ctx_g(cc, 32) != TRANSCRIBE_OK) { @@ -1704,7 +1638,7 @@ transcribe_status run_batch( total_steps > 0 ? step_us / 1000.0 / total_steps : 0.0); } - // ---- Capture results ---- + // Capture results. const int valid_count = std::max(1, static_cast( std::count(valid.begin(), valid.end(), char(1)))); for (int b = 0; b < n; ++b) { diff --git a/src/arch/granite/projector.cpp b/src/arch/granite/projector.cpp index cac2a1e2..9c997c9c 100644 --- a/src/arch/granite/projector.cpp +++ b/src/arch/granite/projector.cpp @@ -3,41 +3,6 @@ // Reference: GraniteSpeechEncoderProjector + // Blip2QFormerLayer/Encoder/Model in transformers/models/blip_2/. // -// One graph end-to-end: -// -// enc_in [hidden_enc, T_enc] + enc_pad [hidden_enc, pad_n] -// → ggml_concat along ne[1] → [hidden_enc, T_pad] -// → ggml_reshape_3d → [hidden_enc, window_size, nblocks] -// query [hidden, num_queries, 1, 1] -// → ggml_repeat_4d → [hidden, num_queries, nblocks, 1] -// for layer in 0..prj_n_layers: -// # Self-attention sublayer (post-LN BERT/BLIP-2 style): -// attn_in = query -// q,k,v = linear(self_attn.{q,k,v}, attn_in) + bias -// attn = softmax(QK^T / sqrt(head_dim)) @ V -// post = linear(self_attn.out, attn) + bias -// query = LN(post + attn_in, norm_self_attn) -// -// # Cross-attention sublayer (cross_attention_frequency=1 ⇒ every layer): -// cross_in = query -// q = linear(cross_attn.q, cross_in) + bias [hidden, num_queries, nblocks] -// k = linear(cross_attn.k, enc_window) + bias [hidden, window_size, nblocks] -// v = linear(cross_attn.v, enc_window) + bias [hidden, window_size, nblocks] -// attn = softmax(QK^T / sqrt(head_dim)) @ V per (head, block) -// post = linear(cross_attn.out, attn) + bias -// query = LN(post + cross_in, norm_cross_attn) -// -// # FFN sublayer: -// ffn_in = query -// mid = GELU(linear(ffn.up, ffn_in) + bias) -// out = linear(ffn.down, mid) + bias -// query = LN(out + ffn_in, norm_ffn) -// -// query = LN(query, qformer.final_norm) # proj.qformer.out [hidden, num_queries, nblocks] -// tokens = reshape(query, hidden, num_queries*nblocks) -// out = linear(proj.linear, tokens) + proj.linear.bias -// # proj.out [text_hidden, num_queries*nblocks] -// // Conventions: // * BLIP-2 layer uses "post-LN" (residual + LN AFTER each output // projection), unlike the granite encoder's pre-LN macaron / pre-LN @@ -98,14 +63,9 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, return y; } -// Compute scaled-dot-product attention with input projections already -// applied. Operates on the Q-Former's standard layout where -// q, k, v all have ne = [hidden, n_seq, batch] -// and the attention runs per `batch` (= nblocks for cross-attn). -// -// Reshape into head form (head_dim = hidden / n_heads), permute so -// `n_seq` ends up at ne[1] for kq via mul_mat(k, q), then softmax over -// k axis, multiply by V^T, and collapse back to ne[hidden, q_seq, batch]. +// Scaled-dot-product attention with input projections already applied. +// q, k, v all have ne = [hidden, n_seq, batch]; attention runs per `batch` +// (= nblocks for cross-attn). head_dim = hidden / n_heads. ggml_tensor * qformer_attn(ggml_context * ctx, ggml_tensor * q_in, // [hidden, q_seq, batch] ggml_tensor * k_in, // [hidden, k_seq, batch] @@ -232,7 +192,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, pb.t_enc_pad = T_pad - T_enc; pb.n_audio_tokens = pb.nblocks * num_queries; - // ----- Graph inputs ----- + // Graph inputs. pb.enc_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, enc_hidden, T_enc); named(pb.enc_in, "proj.enc_in"); ggml_set_input(pb.enc_in); @@ -244,7 +204,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, ggml_set_input(pb.enc_pad); } - // ----- Build window tensor ----- + // Build window tensor. // enc_full: [hidden_enc, T_pad] (concat of enc_in and enc_pad). ggml_tensor * enc_full = pb.enc_in; if (pb.enc_pad != nullptr) { @@ -260,7 +220,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, enc_hidden, window_size, pb.nblocks); - // ----- Build query tensor ----- + // Build query tensor. // weights.proj_top.query is stored at the GGUF's reference dtype // (BF16 for granite). The QFormer layer adds it into the F32 // post-attention activation and feeds the result into LayerNorm; @@ -269,17 +229,11 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, ggml_tensor * query = ggml_cast(ctx, weights.proj_top.query, GGML_TYPE_F32); // Apply the Q-Former INPUT layernorm. Despite the GGUF tensor name - // "proj.qformer.final_norm" (a converter-side misnomer carried over - // from the dump naming), this layernorm corresponds to + // "proj.qformer.final_norm" (a converter-side misnomer), this is // Blip2QFormerModel.layernorm and runs BEFORE the encoder stack: - // - // query_embeds = query_embeds.to(self.layernorm.weight.dtype) // embedding_output = self.layernorm(query_embeds) - // - // Applying it after the layers (as if it were a final norm) silently - // produces wildly drifting output — the layers operate on - // unnormalised queries and the first cross-attention dominates the - // downstream FFN/residual path. + // Applying it after the layers instead silently corrupts the output — + // the layers would operate on unnormalised queries. query = qformer_layer_norm(ctx, query, weights.proj_top.qformer_final_norm_w, weights.proj_top.qformer_final_norm_b, @@ -291,24 +245,21 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, // ggml_repeat_4d emits a fresh contiguous tensor; we can pass it // directly into subsequent ops. - // ----- Q-Former layers ----- + // Q-Former layers. for (int i = 0; i < hp.prj_n_layers; ++i) { query = qformer_layer(ctx, query, enc_window, weights.proj_blocks[i], n_heads, head_dim, ln_eps); } - // ----- proj.qformer.out tap ----- - // The stage-2 dump hooks `proj.qformer` (= Blip2QFormerModel) and - // captures `BaseModelOutputWithPoolingAndCrossAttentions.last_hidden_state`, - // which is the output of the encoder stack — no additional layernorm - // after the layers. Our Q-Former input norm (above) is the only - // wrapper layernorm in the model. + // proj.qformer.out = Blip2QFormerModel.last_hidden_state, the encoder + // stack output — there is no layernorm after the layers (the input norm + // above is the only wrapper layernorm). named(query, "proj.qformer.out"); pb.dumps.qformer_out = query; transcribe::debug::mark_tensor_for_dump(query); - // ----- Reshape to [hidden, num_queries*nblocks] and linear lift ----- + // Reshape to [hidden, num_queries*nblocks] and linear lift. // The reference's .view(batch_size, nblocks*num_queries, hidden) is // a flat reshape over (block_idx, query_idx). In ggml ne layout // (hidden, num_queries, nblocks) the natural reshape to diff --git a/src/arch/granite/projector.h b/src/arch/granite/projector.h index 90f5bab1..a73b9ee5 100644 --- a/src/arch/granite/projector.h +++ b/src/arch/granite/projector.h @@ -2,28 +2,7 @@ // // Reference: GraniteSpeechEncoderProjector + // Blip2QFormerModel/Encoder/Layer in transformers/models/blip_2/. -// -// Forward shape: -// -// in : [hidden_enc, T_enc] (encoder.out for 1b/2b; for -plus -// this is the cat of block[3] and block[N-1] hidden along -// the channel axis, doubling hidden_enc) -// pad : zero-pad along T_enc → nblocks * window_size where -// nblocks = ceil(T_enc / window_size). -// window : view as [hidden_enc, window_size, nblocks] -// query : [hidden, num_queries=3, 1] broadcast → [hidden, 3, nblocks] -// for layer in [0, n_layers): -// q = LN(self_attn.out(self_attn(q, q, q)) + q) -// q = LN(cross_attn.out(cross_attn(q, window, window)) + q) -// q = LN(ffn.down(GELU(ffn.up(q))) + q) -// q = LN(final_norm)(q) # proj.qformer.out [hidden, 3, nblocks] -// out = linear(q.reshape(nblocks*3, hidden)) # proj.out [text_hidden, nblocks*3] -// -// BERT/BLIP-2 layer style: residual + LN happens AFTER each output -// projection ("post-LN" — different from the granite encoder's pre-LN -// macaron / pre-LN attention pattern). -// -// This header is INTERNAL to src/arch/granite/. +// See projector.cpp for the forward conventions. INTERNAL to src/arch/granite/. #pragma once diff --git a/src/arch/granite/weights.cpp b/src/arch/granite/weights.cpp index 1b8c0521..4557e216 100644 --- a/src/arch/granite/weights.cpp +++ b/src/arch/granite/weights.cpp @@ -29,12 +29,12 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, return TRANSCRIBE_ERR_INVALID_ARG; } - // ----- Identity ----- + // Identity. if (auto st = read_optional_string_kv( gguf, "stt.variant", kFamilyTag, "", hp.variant); st != TRANSCRIBE_OK) return st; - // ----- Encoder (Conformer) ----- + // Encoder (Conformer). if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.n_layers", kFamilyTag, hp.enc_n_layers); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.hidden", kFamilyTag, hp.enc_hidden); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.encoder.n_heads", kFamilyTag, hp.enc_n_heads); st != TRANSCRIBE_OK) return st; @@ -63,7 +63,7 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, } } - // ----- Projector (BLIP-2 Q-Former) ----- + // Projector (BLIP-2 Q-Former). if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.n_layers", kFamilyTag, hp.prj_n_layers); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.hidden", kFamilyTag, hp.prj_hidden); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.intermediate", kFamilyTag, hp.prj_intermediate); st != TRANSCRIBE_OK) return st; @@ -75,7 +75,7 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, if (auto st = read_required_u32_kv(gguf, "stt.granite.projector.max_pos_emb", kFamilyTag, hp.prj_max_pos_emb); st != TRANSCRIBE_OK) return st; if (auto st = read_required_string_kv(gguf, "stt.granite.projector.position_embedding_type", kFamilyTag, hp.prj_pos_embed_type); st != TRANSCRIBE_OK) return st; - // ----- Text LM (Granite-4) ----- + // Text LM (Granite-4). if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.n_layers", kFamilyTag, hp.dec_n_layers); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.hidden_size", kFamilyTag, hp.dec_hidden); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.decoder.intermediate_size", kFamilyTag, hp.dec_intermediate); st != TRANSCRIBE_OK) return st; @@ -96,12 +96,12 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.attention_multiplier", kFamilyTag, hp.dec_attention_multiplier); st != TRANSCRIBE_OK) return st; if (auto st = read_required_f32_kv(gguf, "stt.granite.decoder.residual_multiplier", kFamilyTag, hp.dec_residual_multiplier); st != TRANSCRIBE_OK) return st; - // ----- Audio fusion ----- + // Audio fusion. if (auto st = read_required_u32_kv(gguf, "stt.granite.audio_token_id", kFamilyTag, hp.audio_token_id); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.downsample_rate", kFamilyTag, hp.downsample_rate); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite.window_size", kFamilyTag, hp.window_size); st != TRANSCRIBE_OK) return st; - // ----- Frontend (torchaudio MelSpectrogram) ----- + // Frontend (torchaudio MelSpectrogram). if (auto st = read_required_string_kv(gguf, "stt.frontend.type", kFamilyTag, hp.fe_type); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.num_mels", kFamilyTag, hp.fe_num_mels); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.sample_rate", kFamilyTag, hp.fe_sample_rate); st != TRANSCRIBE_OK) return st; @@ -118,7 +118,7 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, if (auto st = read_optional_string_kv(gguf, "stt.frontend.mel_norm", kFamilyTag, "htk", hp.fe_mel_norm); st != TRANSCRIBE_OK) return st; if (auto st = read_optional_bool_kv(gguf, "stt.frontend.center", kFamilyTag, true, hp.fe_center); st != TRANSCRIBE_OK) return st; - // ----- Cross-field invariants ----- + // Cross-field invariants. if (hp.enc_n_layers <= 0 || hp.enc_hidden <= 0 || hp.enc_n_heads <= 0 || hp.enc_head_dim <= 0 || hp.enc_input_dim <= 0 || hp.enc_output_dim <= 0 || hp.enc_feedforward_mult <= 0 || hp.enc_conv_kernel_size <= 0 || @@ -253,9 +253,7 @@ transcribe_status read_granite_hparams(const gguf_context * gguf, return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Weights -// --------------------------------------------------------------------------- +// Weights. namespace { @@ -301,7 +299,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, return TRANSCRIBE_ERR_INVALID_ARG; } - // ----- Derived dims ----- + // Derived dims. const int64_t enc_h = hp.enc_hidden; const int64_t enc_in = hp.enc_input_dim; const int64_t enc_out = hp.enc_output_dim; @@ -330,7 +328,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, const int64_t q_out = dec_nh * dec_hd; const int64_t kv_out = dec_nkv * dec_hd; - // ----- Encoder: top-level ----- + // Encoder: top-level. GET_LIN(weights.enc_top.input_linear_w, "enc.input_linear.weight", enc_in, enc_h); GET_F32(weights.enc_top.input_linear_b, "enc.input_linear.bias", enc_h); GET_LIN(weights.enc_top.ctc_proj_w, "enc.ctc_proj.weight", enc_h, enc_out); @@ -338,7 +336,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, GET_LIN(weights.enc_top.ctc_bypass_w, "enc.ctc_bypass.weight", enc_out, enc_h); GET_F32(weights.enc_top.ctc_bypass_b, "enc.ctc_bypass.bias", enc_h); - // ----- Encoder: per-block ----- + // Encoder: per-block. weights.enc_blocks.assign(hp.enc_n_layers, GraniteEncBlock{}); for (int i = 0; i < hp.enc_n_layers; ++i) { auto & b = weights.enc_blocks[i]; @@ -402,7 +400,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, GET_F32(b.norm_post_b, lname("enc.blocks.%d.norm_post.bias", i), enc_h); } - // ----- Projector: top-level ----- + // Projector: top-level. // // proj.query has shape (1, num_queries, hidden) in PyTorch, which // is ne[hidden, num_queries, 1] in ggml. num_queries is not in KV; @@ -421,7 +419,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, GET_F32(weights.proj_top.qformer_final_norm_b, "proj.qformer.final_norm.bias", prj_h); - // ----- Projector: per-Q-Former-layer ----- + // Projector: per-Q-Former-layer. weights.proj_blocks.assign(prj_n_layers, GraniteProjBlock{}); for (int i = 0; i < prj_n_layers; ++i) { auto & b = weights.proj_blocks[i]; @@ -461,10 +459,10 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, GET_F32(b.norm_ffn_b, lname("proj.qformer.blocks.%d.norm_ffn.bias", i), prj_h); } - // ----- Text LM: embed ----- + // Text LM: embed. GET_LIN(weights.dec_embed.token_w, "dec.token_embd.weight", dec_h, dec_vocab); - // ----- Text LM: blocks ----- + // Text LM: blocks. weights.dec_blocks.assign(hp.dec_n_layers, GraniteDecBlock{}); for (int i = 0; i < hp.dec_n_layers; ++i) { auto & b = weights.dec_blocks[i]; @@ -485,7 +483,7 @@ transcribe_status build_granite_weights(ggml_context * ctx_meta, } } - // ----- Text LM: final norm + (conditional) output ----- + // Text LM: final norm + (conditional) output. GET_F32(weights.dec_final.norm_w, "dec.output_norm.weight", dec_h); if (hp.dec_tie_word_embeddings) { // Tied head (granite-speech-4.1-2b-plus): no dec.output.weight diff --git a/src/arch/granite/weights.h b/src/arch/granite/weights.h index bc771d08..c1addae8 100644 --- a/src/arch/granite/weights.h +++ b/src/arch/granite/weights.h @@ -41,9 +41,7 @@ struct ggml_tensor; namespace transcribe::granite { -// --------------------------------------------------------------------------- // Hyperparameters -// --------------------------------------------------------------------------- struct GraniteHParams { // Identity. `stt.variant` distinguishes models that share the same @@ -119,8 +117,7 @@ struct GraniteHParams { int32_t vocab_size = 0; // 100353 // Frontend (torchaudio MelSpectrogram). htk-norm mel, no log, no - // per-utterance norm, power=2 spectrogram. Stage 2b extends - // MelFrontend to support these knobs; for now we just record them. + // per-utterance norm, power=2 spectrogram. std::string fe_type; // "mel" int32_t fe_num_mels = 0; // 80 int32_t fe_sample_rate = 0; // 16000 @@ -141,9 +138,7 @@ struct GraniteHParams { transcribe_status read_granite_hparams(const gguf_context * gguf, GraniteHParams & hp); -// --------------------------------------------------------------------------- // Weight slots - encoder (Conformer) -// --------------------------------------------------------------------------- // Top-level encoder tensors: input linear (160 -> 1024) and the // self-conditioned CTC bypass pair (ctc_proj: 1024 -> 348 for the @@ -183,9 +178,7 @@ struct GraniteEncBlock { // Looked up by clamped (k - q + max_pos_emb) and added to QK^T. ggml_tensor * attn_rel_pos_emb = nullptr; // [head_dim, 2*max_pos_emb+1] - // Conv module: LN -> pointwise expand to 2*inner_dim -> GLU split - // -> depthwise(k=15) -> BN -> SiLU -> pointwise contract back to - // hidden. inner_dim = hidden * conv_expansion. + // Conv module (see granite_conv_module in encoder.cpp). ggml_tensor * norm_conv_w = nullptr; // [hidden] ggml_tensor * norm_conv_b = nullptr; ggml_tensor * conv_pointwise1_w = nullptr; // [1, hidden, 2*inner_dim] @@ -198,10 +191,7 @@ struct GraniteEncBlock { ggml_tensor * conv_pointwise2_w = nullptr; // [1, inner_dim, hidden] ggml_tensor * conv_pointwise2_b = nullptr; // [hidden] // Fused BatchNorm scale/bias precomputed at load (model.cpp - // fuse_batch_norm). Replaces the raw weight/bias + running stats - // for graph use; the raw tensors are still loaded but only feed - // the fusion step. Both are [inner_dim] f32, allocated in a - // separate ctx on the model. + // fuse_batch_norm); the raw tensors above only feed the fusion step. ggml_tensor * conv_bn_fused_scale = nullptr; // [inner_dim] ggml_tensor * conv_bn_fused_bias = nullptr; // [inner_dim] @@ -218,15 +208,16 @@ struct GraniteEncBlock { ggml_tensor * norm_post_b = nullptr; }; -// --------------------------------------------------------------------------- // Weight slots - projector (BLIP-2 Q-Former) -// --------------------------------------------------------------------------- // Top-level projector tensors. // query: [hidden, num_queries=3, 1] — learned queries broadcast over // per-window encoder slices. // linear: lifts Q-Former hidden (1024) to LM hidden (2048). -// qformer.final_norm: post Q-Former-stack LN over the query stream. +// qformer.final_norm: despite the name (a converter-side misnomer), +// this is Blip2QFormerModel.layernorm — the Q-Former INPUT LN, +// applied to the queries BEFORE the layer stack, not after it. +// See projector.cpp (qformer_layer_norm at graph entry). struct GraniteProjTop { ggml_tensor * query = nullptr; // [hidden, num_queries, 1] ggml_tensor * linear_w = nullptr; // [hidden, lm_hidden] @@ -273,9 +264,7 @@ struct GraniteProjBlock { ggml_tensor * norm_ffn_b = nullptr; }; -// --------------------------------------------------------------------------- // Weight slots - Granite-4 LM -// --------------------------------------------------------------------------- struct GraniteDecEmbed { ggml_tensor * token_w = nullptr; // [hidden, vocab] @@ -305,9 +294,7 @@ struct GraniteDecFinal { ggml_tensor * output_w = nullptr; // [hidden, vocab] or nullptr }; -// --------------------------------------------------------------------------- // Aggregated weights -// --------------------------------------------------------------------------- struct GraniteWeights { GraniteEncTop enc_top; diff --git a/src/arch/granite_nar/decoder.cpp b/src/arch/granite_nar/decoder.cpp index 0b1d7a9b..a823247e 100644 --- a/src/arch/granite_nar/decoder.cpp +++ b/src/arch/granite_nar/decoder.cpp @@ -177,7 +177,7 @@ ForwardBuild build_forward_graph(ggml_context * ctx, const float emb_mul = hp.dec_embedding_multiplier; const auto params = to_params(hp); - // ---------- Inputs ---------- + // Inputs. fb.audio_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, n_audio_tokens); named(fb.audio_in, "dec.audio_in"); ggml_set_input(fb.audio_in); @@ -197,7 +197,7 @@ ForwardBuild build_forward_graph(ggml_context * ctx, } fb.graph = gf; - // ---------- Flat input embeds ---------- + // Flat input embeds. // text embeds via lookup. ggml_get_rows returns the embed table dtype // (BF16 here); cast to F32 so ggml_concat with the F32 audio_in is // legal (concat requires matching dtypes). @@ -221,16 +221,16 @@ ForwardBuild build_forward_graph(ggml_context * ctx, // round-trip through ×(1/emb_mul)×emb_mul = identity). ggml_tensor * x = ggml_scale(ctx, flat, emb_mul); - // ---------- Block stack ---------- + // Block stack. for (int il = 0; il < n_layer; ++il) { x = block_bidi(ctx, x, weights.dec_blocks[il], params, fb.T_total, fb.positions_in); } - // ---------- Final RMSNorm ---------- + // Final RMSNorm. x = rms_norm(ctx, x, weights.dec_final.norm_w, rms_eps); - // ---------- Slice text portion ---------- + // Slice text portion. // text_x = x[:, n_audio_tokens:] ggml_tensor * text_x = ggml_view_2d( ctx, x, @@ -240,18 +240,10 @@ ForwardBuild build_forward_graph(ggml_context * ctx, static_cast(n_audio_tokens)); text_x = ggml_cont(ctx, text_x); - // ---------- LM head (tied to token_embd) ---------- - // 99a4df9 modeling: GraniteSpeechNarLM.forward applies - // logits = lm_head(hidden) / config.logits_scaling - // (the /logits_scaling divisor is part of the Granite-4 head and - // the new NAR LM does NOT bypass it — see line 913 of - // modeling_granite_speech_nar.py). The OLD multi-file snapshot - // (7d20732d) called self.llm.lm_head(text_x) directly, bypassing - // the divisor. The dumper now follows the README path which - // routes through the new forward; mirror its scaling here so - // dec.text_logits dump aligns numerically with the reference. - // Argmax is scale-invariant either way, so the C++ transcript is - // unchanged by this division. + // LM head (tied to token_embd). The Granite-4 head divides by + // logits_scaling; we mirror it so dec.text_logits aligns numerically with + // the reference dump. Argmax is scale-invariant, so the transcript is + // unchanged by this division. ggml_tensor * logits = ggml_mul_mat( ctx, weights.dec_embed.token_w, text_x); if (hp.dec_logits_scaling > 0.0f && hp.dec_logits_scaling != 1.0f) { diff --git a/src/arch/granite_nar/encoder.cpp b/src/arch/granite_nar/encoder.cpp index c93d9a62..f20272f8 100644 --- a/src/arch/granite_nar/encoder.cpp +++ b/src/arch/granite_nar/encoder.cpp @@ -4,7 +4,7 @@ // 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 (1024 → 100353). We emit +// - 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. @@ -64,9 +64,7 @@ ggml_tensor * linear(ggml_context * ctx, ggml_tensor * x, } // namespace -// --------------------------------------------------------------------------- -// Host-side mel + 2-frame stack -// --------------------------------------------------------------------------- +// Host-side mel + 2-frame stack. transcribe_status compute_mel_encoder_input( const transcribe::MelFrontend & mel, @@ -121,9 +119,7 @@ transcribe_status compute_mel_encoder_input( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Shaw bookkeeping (matches AR granite) -// --------------------------------------------------------------------------- +// Shaw bookkeeping (matches AR granite). std::vector precompute_attention_dists(int context_size, int max_pos_emb) { @@ -155,10 +151,8 @@ std::vector precompute_last_block_mask(int context_size, int t_enc_remain return mask; } -// --------------------------------------------------------------------------- // Per-block builders (identical to AR granite, parameterised over the // granite_nar weight struct). -// --------------------------------------------------------------------------- namespace { @@ -297,9 +291,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const auto & b = weights.enc_blocks[i]; x = macaron(ctx, x, b, /*is_ff1=*/true); - // Sub-step dump: post-FF1 residual on block 0. Used by - // validate.py to localize bf16-cascade drift within the - // first conformer block. + // Sub-step dump: post-FF1 residual on block 0. if (i == 0) { named(x, "enc.block.0.post_ff1"); eb.dumps.block_0_post_ff1 = x; @@ -472,9 +464,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } -// --------------------------------------------------------------------------- -// Host-side BPE CTC pool + greedy decode -// --------------------------------------------------------------------------- +// Host-side BPE CTC pool + greedy decode. void compute_bpe_ctc_initial_hypothesis( const std::vector & importance_non_blank, @@ -536,16 +526,12 @@ void compute_bpe_ctc_initial_hypothesis( } } if (argmax != blank_id && argmax != prev) { - // Two BPE-CTC schemes are in flight depending on the GGUF's - // source snapshot: - // - Old (bpe_output_dim = vocab_size + 1, blank_id = 0): - // channel 0 is a synthetic blank, channels 1..N hold the - // LLM token ids — we recover the LLM id with `argmax - 1`. - // - New (bpe_output_dim = vocab_size, blank_id = BOS=100257): - // channels ARE the LLM ids directly; the blank channel IS - // one of them (the BOS id). No shift needed. - // The two are distinguished by blank_id alone: id 0 means the - // old scheme, anything else means the new direct-id scheme. + // 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); } diff --git a/src/arch/granite_nar/encoder.h b/src/arch/granite_nar/encoder.h index 6ce400b6..36e29b53 100644 --- a/src/arch/granite_nar/encoder.h +++ b/src/arch/granite_nar/encoder.h @@ -57,7 +57,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_out_dim=100353, 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]. // Used host-side as the BPE pool's @@ -75,7 +75,7 @@ struct EncoderBuild { // (== `enc.block.8.out`) ggml_tensor * block_last_out = nullptr; // block (N-1) ggml_tensor * ctc_logits = nullptr; // == ctc_logits, named - // Block-0 sub-step taps for bf16-cascade drift localization. + // Block-0 sub-step taps. ggml_tensor * block_0_post_ff1 = nullptr; ggml_tensor * block_0_post_attn = nullptr; ggml_tensor * block_0_post_conv = nullptr; @@ -103,26 +103,26 @@ std::vector precompute_last_block_mask(int context_size, int t_enc_rema // 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, 100353] host-row-major (frame-level BPE logits) +// 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 0 +// 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: argmax over the -// CTC head per frame, only frames where the argmax is non-blank are -// "valid" — those frames are then bucketed into windows of pool_window=4 -// and the BPE logits over each window are weighted by the (softmaxed) -// CTC blank-vs-non-blank posterior. We approximate the same path: -// -// per-frame: blank_prob = softmax(ctc_logits)[blank_id] -// non_blank_prob = 1 - blank_prob -// pool: for each window of pool_window consecutive frames whose -// non_blank_prob > 0.5, sum non_blank_prob[k] * ctc_bpe[k] -// normalised by sum non_blank_prob[k] -> [bpe_dim] -// greedy: argmax per pooled window, then collapse repeats + drop blanks +// 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, diff --git a/src/arch/granite_nar/model.cpp b/src/arch/granite_nar/model.cpp index 71527183..7f0efc07 100644 --- a/src/arch/granite_nar/model.cpp +++ b/src/arch/granite_nar/model.cpp @@ -2,7 +2,7 @@ // // Architecture: // - Conformer encoder (16 layers, hidden=1024) with self-conditioned -// CTC bypass at layer N/2 and a BPE-CTC head (1024 -> 100353) +// CTC bypass at layer N/2 and a BPE-CTC head // - EncoderProjectorQFormer (per-encoder-layer LN over the cat of 4 // captured layers, layer_projector 4096->2048, learned query + // window_positions, 2 cross-attn+MLP layers, out_norm + out_linear) @@ -98,7 +98,6 @@ namespace { constexpr const char k_default_variant[] = "granite-speech-nar"; constexpr float kBnEps = 1e-5f; -// --------------------------------------------------------------------------- // Input-length contract (see docs/input-limits.md). granite_nar is a // hard-context-cap family: the bidirectional editor runs a SINGLE forward // pass over (audio_tokens + text) and every position uses RoPE bounded by @@ -108,7 +107,6 @@ constexpr float kBnEps = 1e-5f; // the up-front input gate on T_total plus an advisory max_audio_ms. Over- // length input is rejected with TRANSCRIBE_ERR_INPUT_TOO_LONG before the // decoder graph is built. -// --------------------------------------------------------------------------- // Representative text budget reserved when advertising max_audio_ms. The // editor's text side (initial BPE hypothesis + insertion slots) shares the @@ -138,15 +136,12 @@ int granite_nar_num_queries(const GraniteNarHParams & hp) { } // Advisory transcribe_capabilities::max_audio_ms: the longest audio whose -// audio tokens, a representative text hypothesis (k_text_budget), still fit +// audio tokens plus a representative text hypothesis (k_text_budget) still fit // the decoder context ceiling. This is the input bound the gate enforces. // Returns 0 ("unknown / unbounded") if the rate constants are missing, so a -// misconfigured model is never advertised with a wrong finite number. -// -// Mirrors granite/model.cpp's granite_max_audio_ms with granite_nar's hparam -// field names (prj_block_size / prj_downsample_rate vs window_size / -// downsample_rate) and reserves a text budget in place of AR's prompt + -// generation reserve, because the NAR editor has no autoregressive output. +// misconfigured model is never advertised with a wrong finite number. Mirrors +// granite's granite_max_audio_ms, reserving a text budget in place of AR's +// prompt + generation reserve. int64_t granite_nar_max_audio_ms(const GraniteNarHParams & hp) { const int num_queries = granite_nar_num_queries(hp); if (num_queries <= 0 || hp.prj_block_size <= 0 || @@ -267,7 +262,7 @@ transcribe_status load( // Publish the input-length ceiling now that the decoder context window and // frontend rate are known (apply_family_invariants ran before the hparams - // were read, so it could not set this). See docs/input-limits.md. + // were read, so it could not set this). m->caps.max_audio_ms = granite_nar_max_audio_ms(m->hparams); // Basis for the session-level limits query (transcribe_session_get_limits): @@ -302,7 +297,7 @@ transcribe_status load( } } - // ---------- Mel frontend ---------- + // Mel frontend. // The NLE shares the whisper-mode feature extractor with AR Granite // but the NAR GGUF doesn't ship pre_emphasis / f_min / f_max KVs (the // upstream WhisperFeatureExtractor doesn't expose them; we hardcode @@ -347,7 +342,7 @@ transcribe_status load( m->mel.emplace(cfg); } - // ---------- Tensor catalog ---------- + // Tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -455,7 +450,7 @@ transcribe_status run( return TRANSCRIBE_ERR_INVALID_ARG; } - // ---------- Mel + 2-frame stack ---------- + // Mel + 2-frame stack. const int64_t t_mel_start = ggml_time_us(); int t_enc = 0; if (const transcribe_status mst = compute_mel_encoder_input( @@ -478,7 +473,7 @@ transcribe_status run( shape, 2, "encoder"); } - // ---------- Reset compute state ---------- + // Reset compute state. if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; @@ -497,7 +492,7 @@ transcribe_status run( } } - // ---------- Encoder graph ---------- + // Encoder graph. EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, t_enc, cc->encoder_use_flash); @@ -583,7 +578,7 @@ transcribe_status run( eb.dumps.block_last_out, "encoder"); try_dump("enc.ctc_logits", eb.dumps.ctc_logits, "encoder"); - // ---------- Read encoder outputs to host ---------- + // Read encoder outputs to host. const int64_t cat_h = eb.cat_out->ne[0]; const int64_t cat_T = eb.cat_out->ne[1]; cc->enc_cat_host.resize(static_cast(cat_h) * cat_T); @@ -623,7 +618,7 @@ transcribe_status run( // and gets the v-th logit at frame t correctly. (void)n_ctc_vocab; - // ---------- Initial BPE hypothesis ---------- + // Initial BPE hypothesis. std::vector hyp_ids; if (n_bpe_vocab > 0 && !mid_non_blank.empty()) { compute_bpe_ctc_initial_hypothesis( @@ -646,7 +641,7 @@ transcribe_status run( return TRANSCRIBE_OK; } - // ---------- Projector graph ---------- + // Projector graph. ggml_context * proj_ctx = nullptr; { ggml_init_params ip {}; @@ -720,12 +715,12 @@ transcribe_status run( cc->proj_out_host.size() * sizeof(float)); ggml_free(proj_ctx); - // ---------- add_insertion_slots ---------- + // add_insertion_slots. std::vector text_ids; add_insertion_slots(hyp_ids, cm->hparams.dec_eos_id, text_ids); const int n_text = static_cast(text_ids.size()); - // ---------- Input-length gate (see docs/input-limits.md) ---------- + // Input-length gate. // The editor is non-autoregressive: a SINGLE bidirectional forward over // (audio_tokens + text) where every position's RoPE is bounded by // dec_max_pos_emb (optionally lowered by the caller's n_ctx knob). Both @@ -756,7 +751,7 @@ transcribe_status run( for (auto & v : cc->proj_out_host) v *= inv; } - // ---------- Decoder graph ---------- + // Decoder graph. ggml_context * dec_ctx = nullptr; { ggml_init_params ip {}; diff --git a/src/arch/granite_nar/projector.cpp b/src/arch/granite_nar/projector.cpp index d2186de1..2e32f1d7 100644 --- a/src/arch/granite_nar/projector.cpp +++ b/src/arch/granite_nar/projector.cpp @@ -119,7 +119,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, named(pb.enc_in, "proj.enc_in"); ggml_set_input(pb.enc_in); - // ---------- Per-encoder-layer LayerNorms ---------- + // Per-encoder-layer LayerNorms. // Slice enc_in along the channel axis into `enc_layers` chunks of // size `enc_hidden`, LN each chunk with its own gamma/beta, then // concat back. @@ -141,14 +141,14 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, } // normed ne = [cat_in, T_enc] - // ---------- Linear projector (cat_in -> prj_hidden) + GELU ---------- + // Linear projector (cat_in -> prj_hidden) + GELU. ggml_tensor * proj_lp = linear(ctx, normed, weights.proj_top.layer_projector_w, weights.proj_top.layer_projector_b); proj_lp = ggml_gelu(ctx, proj_lp); // proj_lp ne = [prj_hidden, T_enc] - // ---------- Pad along T to nblocks * block_size ---------- + // Pad along T to nblocks * block_size. // The pad is allocated at prj_hidden width (post-GELU) so the // caller only uploads a small zero buffer once per encode. The // reference does the pad on the same axis, post-layer_projector, @@ -167,7 +167,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, ggml_tensor * windowed = ggml_reshape_3d(ctx, full, prj_hidden, block_size, pb.nblocks); - // ---------- K/V: x + window_positions ---------- + // K/V: x + window_positions. // window_positions tensor ne = [prj_hidden, block_size, 1]. We add // it to `windowed` which is [prj_hidden, block_size, nblocks]; ggml // broadcasts the ne[2]=1 across the nblocks dim. @@ -175,7 +175,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, GGML_TYPE_F32); ggml_tensor * kv = ggml_add(ctx, windowed, wpos); - // ---------- Mean-pool windowed -> [prj_hidden, n_query, nblocks] ---------- + // Mean-pool windowed -> [prj_hidden, n_query, nblocks]. // Reference reshapes the windowed tensor (B*nblocks, block_size, h) as // (B*nblocks, n_query, downsample, h) and means over the downsample // axis. In ggml ne order, this is reshape [h, block_size, nblocks] -> @@ -192,14 +192,14 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, ggml_tensor * mean_pool = ggml_reshape_3d(ctx, mp_mean, prj_hidden, n_query, pb.nblocks); - // ---------- Query: learned [prj_hidden, n_query, 1] + mean_pool ---------- + // Query: learned [prj_hidden, n_query, 1] + mean_pool. ggml_tensor * query = ggml_cast(ctx, weights.proj_top.query, GGML_TYPE_F32); // Broadcast query over the nblocks axis when adding mean_pool. ggml_add // broadcasts dims of size 1, so the [h, n_query, 1] query expands // naturally to match mean_pool's [h, n_query, nblocks]. query = ggml_add(ctx, mean_pool, query); - // ---------- Q-Former layers (no self-attention) ---------- + // Q-Former layers (no self-attention). for (int i = 0; i < hp.prj_n_layers; ++i) { const auto & b = weights.proj_blocks[i]; @@ -226,7 +226,7 @@ ProjectorBuild build_projector_graph(ggml_context * ctx, pb.dumps.qformer_out = query; transcribe::debug::mark_tensor_for_dump(query); - // ---------- out_norm + out_linear (LLM-space) ---------- + // out_norm + out_linear (LLM-space). ggml_tensor * onorm = layer_norm(ctx, query, weights.proj_top.out_norm_w, weights.proj_top.out_norm_b, diff --git a/src/arch/granite_nar/projector.h b/src/arch/granite_nar/projector.h index ac9d78d8..67e3337e 100644 --- a/src/arch/granite_nar/projector.h +++ b/src/arch/granite_nar/projector.h @@ -6,37 +6,11 @@ // the encoder K/V plus a learned `window_positions` bias acts in place // of positional encodings. // -// Forward shape (per upstream): +// The upstream layer math is PRE-LN (not post-LN like the AR BLIP-2 layer): +// norm comes BEFORE each sublayer's linear cluster, and the residual closes +// after the second linear. See projector.cpp for the forward graph. // -// in [4 * 1024, T_enc] // 4 captured encoder layers concatenated -// for j in [0..4): -// in[j*1024:(j+1)*1024] = per_layer_LN_j(in[j*1024:(j+1)*1024]) -// in = layer_projector(in) + GELU // [2048, T_enc] -// pad along T_enc → nblocks * 15 -// reshape to [2048, 15, nblocks] -// kv = mean over the window axis -> [2048, nblocks] -// kv = kv + window_positions (broadcast across nblocks) // [2048, 15, nblocks] -// query = learned [2048, 3, 1] repeated across nblocks -> [2048, 3, nblocks] -// for layer in [0, prj_n_layers): -// # cross-attention sublayer (post-LN, BLIP-2 style): -// q = norm_attn(query) -// q' = linear(cross_attn_q, q) -// k' = linear(cross_attn_k, kv) -// v' = linear(cross_attn_v, kv) -// attn = softmax(q' k'^T / sqrt(head_dim)) v' -// query = query + linear(cross_attn_o, attn) -// # FFN sublayer: -// ff_in = norm_ffn(query) -// query = query + linear(ffn_fc2, SiLU(linear(ffn_fc1, ff_in))) -// query = out_norm(query) -// tokens = reshape(query, [2048, 3 * nblocks]) -// out = linear(out_linear, tokens) // [2048, 3 * nblocks] -// -// Note that the upstream layer math is PRE-LN (not post-LN like the AR -// BLIP-2 layer): norm comes BEFORE each sublayer's linear cluster, and -// the residual closes after the second linear. -// -// This header is INTERNAL to src/arch/granite_nar/. +// INTERNAL to src/arch/granite_nar/. #pragma once diff --git a/src/arch/granite_nar/weights.cpp b/src/arch/granite_nar/weights.cpp index e9b9ac75..73de9e67 100644 --- a/src/arch/granite_nar/weights.cpp +++ b/src/arch/granite_nar/weights.cpp @@ -60,9 +60,8 @@ transcribe_status read_granite_nar_hparams(const gguf_context * gguf, if (auto st = read_required_u32_kv(gguf, "stt.granite_nar.encoder.output_dim", kTag, hp.enc_output_dim); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite_nar.encoder.bpe_output_dim", kTag, hp.enc_bpe_output_dim); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.granite_nar.encoder.bpe_pool_window", kTag, hp.enc_bpe_pool_window); st != TRANSCRIBE_OK) return st; - // bpe_blank_id is a NEW key (added when we re-finalized against the - // 99a4df9 snapshot). Older GGUFs (converted before this change) won't - // have it; default to 0 to preserve the legacy decode scheme. + // bpe_blank_id is an optional key; older GGUFs lack it. Default to 0 to + // preserve the legacy decode scheme. { int32_t tmp = 0; if (read_required_u32_kv(gguf, "stt.granite_nar.encoder.bpe_blank_id", kTag, tmp) == TRANSCRIBE_OK) { @@ -206,7 +205,7 @@ transcribe_status build_granite_nar_weights(ggml_context * ctx_meta, const int64_t q_out = dec_nh * dec_hd; const int64_t kv_out = dec_nkv * dec_hd; - // ----- Encoder top ----- + // Encoder top. GET_LIN(weights.enc_top.input_linear_w, "enc.input_linear.weight", enc_in, enc_h); GET_F32(weights.enc_top.input_linear_b, "enc.input_linear.bias", enc_h); GET_LIN(weights.enc_top.ctc_proj_w, "enc.ctc_proj.weight", enc_h, enc_out); @@ -218,7 +217,7 @@ transcribe_status build_granite_nar_weights(ggml_context * ctx_meta, GET_F32(weights.enc_top.ctc_bpe_b, "enc.ctc_bpe.bias", enc_bpe_out); } - // ----- Encoder blocks ----- + // Encoder blocks. weights.enc_blocks.assign(hp.enc_n_layers, GraniteNarEncBlock{}); for (int i = 0; i < hp.enc_n_layers; ++i) { auto & b = weights.enc_blocks[i]; @@ -262,7 +261,7 @@ transcribe_status build_granite_nar_weights(ggml_context * ctx_meta, GET_F32(b.norm_post_b, lname("enc.blocks.%d.norm_post.bias", i), enc_h); } - // ----- Projector top ----- + // Projector top. GET_LIN(weights.proj_top.layer_projector_w, "prj.layer_projector.weight", prj_kv_in, prj_h); GET_F32(weights.proj_top.layer_projector_b, "prj.layer_projector.bias", prj_h); GET_F32(weights.proj_top.out_norm_w, "prj.out_norm.weight", prj_h); @@ -320,7 +319,7 @@ transcribe_status build_granite_nar_weights(ggml_context * ctx_meta, GET_F32(b.ffn_fc2_b, lname("prj.blocks.%d.ffn_fc2.bias", i), prj_h); } - // ----- LLM (granite-4) ----- + // LLM (granite-4). GET_LIN(weights.dec_embed.token_w, "dec.token_embd.weight", dec_h, dec_vocab); weights.dec_blocks.assign(hp.dec_n_layers, GraniteNarDecBlock{}); diff --git a/src/arch/granite_nar/weights.h b/src/arch/granite_nar/weights.h index c8e338a1..8a7becf4 100644 --- a/src/arch/granite_nar/weights.h +++ b/src/arch/granite_nar/weights.h @@ -8,7 +8,7 @@ // // enc.* Conformer encoder (16 layers, hidden=1024) — same block as // the AR granite encoder, plus an extra `enc.ctc_bpe.*` head -// (1024 -> 100353) used for the initial BPE CTC hypothesis. +// used for the initial BPE CTC hypothesis. // prj.* EncoderProjectorQFormer: // - 4 per-encoder-layer LayerNorms (one per layer index in // encoder_layer_indices, e.g. [4, 8, 12, -1]) @@ -36,9 +36,7 @@ struct ggml_tensor; namespace transcribe::granite_nar { -// --------------------------------------------------------------------------- // Hyperparameters -// --------------------------------------------------------------------------- struct GraniteNarHParams { // Encoder @@ -130,9 +128,7 @@ struct GraniteNarHParams { transcribe_status read_granite_nar_hparams(const gguf_context * gguf, GraniteNarHParams & hp); -// --------------------------------------------------------------------------- // Weight slots -// --------------------------------------------------------------------------- struct GraniteNarEncTop { ggml_tensor * input_linear_w = nullptr; @@ -141,7 +137,7 @@ struct GraniteNarEncTop { ggml_tensor * ctc_proj_b = nullptr; ggml_tensor * ctc_bypass_w = nullptr; // 348 -> 1024 (out_mid) ggml_tensor * ctc_bypass_b = nullptr; - ggml_tensor * ctc_bpe_w = nullptr; // 1024 -> 100353 (NLE-only) + ggml_tensor * ctc_bpe_w = nullptr; // 1024 -> bpe_output_dim ggml_tensor * ctc_bpe_b = nullptr; }; diff --git a/src/arch/medasr/encoder.cpp b/src/arch/medasr/encoder.cpp index d9ecffef..858f911f 100644 --- a/src/arch/medasr/encoder.cpp +++ b/src/arch/medasr/encoder.cpp @@ -1,39 +1,13 @@ // arch/medasr/encoder.cpp - MedASR (Google LASR-CTC) Conformer encoder // graph builder. // -// Forward shape: -// -// mel.in ne=[T_mel, n_mels=128, 1, B] -// -> subsampling (custom: dense + 2x conv1d(s=2,p=0) + dense) -// ne=[d_model=512, T_enc, 1, B] = enc.subsampling.out -// -> 17 x conformer block { -// FF1 macaron residual w=[1.5, 0.5] (sub-step: enc.block.0.post_ff1) -// RoPE self-attention, unscaled residual (post_attn) -// conv module, scaled residual w=[2.0, 1.0] (post_conv) -// FF2 macaron residual w=[1.5, 0.5] (post_ff2) -// norm_out = enc.block.i.out -// } -// -> enc.out_norm = enc.out_norm.out -// -> ctc head (Conv1d k=1 = mul_mat + bias) -// ne=[vocab=512, T_enc, 1, B] = enc.ctc_logits -// -// Differences from src/arch/gigaam/encoder.cpp (the closest in-tree -// analog): -// -// - Subsampling is 1-D (linear -> relu -> conv1d(s=2) -> relu -> -// conv1d(s=2) -> relu -> linear), NOT 2x conv1d. The leading dense -// reduces n_mels -> d_model before the convs touch the time axis. -// - All encoder LayerNorms are bias=false (norm has scale w, no beta). -// - Conv module uses BatchNorm — host-fused at load time, uploaded as -// compute_ctx tensors per call. Asymmetric padding (15, 16) for -// k=32 "same" depthwise. -// - Scaled residuals: macaron weights [1.5, 0.5] on FF1/FF2, conv -// weights [2.0, 1.0]. NOT the standard 0.5 macaron half-step. -// - Attention uses RoPE on Q and K AFTER projection (standard -// convention). GigaAM rotates before projection, which is rare. -// - LayerNorm eps is 1e-6 (loaded from hparams), NOT the -// conf::layer_norm hardcoded 1e-5; this file uses a local helper -// `lasr_layer_norm` that takes eps as a runtime arg. +// mel.in ne=[n_mels=128, T_mel, 1, B] -> subsampling (dense + 2x +// conv1d(s=2,p=0) + dense) [d_model=512, T_enc, 1, B] -> 17 x conformer +// block {macaron FF1 [1.5,0.5] | RoPE self-attn | conv (BatchNorm) [2.0,1.0] +// | macaron FF2 [1.5,0.5] | norm_out} -> enc.out_norm -> ctc head +// (Conv1d k=1) [vocab=512, T_enc, 1, B]. All encoder LayerNorms are +// bias=false with eps=1e-6 (local `lasr_layer_norm`, not conf's 1e-5); +// RoPE is applied to Q/K after projection. #include "encoder.h" @@ -56,17 +30,13 @@ namespace { namespace conf = transcribe::conformer; -// Weight-matmul forcing F32 accumulation for F16 weights. On CUDA, -// `ggml_mul_mat` with an F16 weight takes the cuBLAS COMPUTE_16F path -// whose accumulator saturates at ~6.5e4 — and medasr's scaled-residual -// stream pushes intermediate activations to ~2e6 between blocks (the -// macaron [1.5, 0.5] and conv [2.0, 1.0] amplification). F16 accum -// overflows to NaN, the CTC argmax collapses to blank/specials, and -// every transcript comes back empty. GGML_PREC_F32 forces cuBLAS to run -// the GEMM in F32 instead, matching CPU/Metal which always F32-accumulate. -// No-op on CPU/Metal; slight perf cost on CUDA. Skip on non-F16 weights -// (BF16 already COMPUTE_32F; quantized routes through MMQ). Mirrors the -// `mul_mat_f32acc` helper at src/causal_lm/causal_lm.cpp:40. +// Weight-matmul forcing F32 accumulation for F16 weights. On CUDA, an F16 +// weight takes the cuBLAS COMPUTE_16F path whose accumulator saturates at +// ~6.5e4 — and medasr's scaled-residual stream (macaron [1.5,0.5] + conv +// [2.0,1.0]) pushes activations to ~2e6 between blocks, so F16 accum +// overflows to NaN and every transcript comes back empty. GGML_PREC_F32 +// matches CPU/Metal's always-F32 accumulate. No-op on CPU/Metal and on +// non-F16 weights (BF16 already COMPUTE_32F; quantized routes through MMQ). ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * x) @@ -122,31 +92,12 @@ ggml_tensor * scaled_residual(ggml_context * ctx, } // --------------------------------------------------------------------------- -// Subsampling +// Subsampling: dense_0 -> relu -> conv_0(s=2) -> relu -> conv_1(s=2) -> +// relu -> dense_1 (no relu). The running tensor stays "channel-fast" +// [C, T, B] for the dense layers (mul_mat) and "time-fast" [T, C, B] for +// the convs (conv_1d_f32); one permute on each boundary. Per-step ne is +// noted inline below. // --------------------------------------------------------------------------- -// -// Reference (LasrEncoderSubsampling.forward): -// h = relu(dense_0(x)) # [B, T_mel, d_model] -// h = h.transpose(1, 2) # [B, d_model, T_mel] -// h = relu(conv_0(h)) # [B, d_model, T1] (k=5, s=2, p=0) -// h = relu(conv_1(h)) # [B, sub_c, T_enc] (k=5, s=2, p=0) -// h = h.transpose(1, 2) # [B, T_enc, sub_c] -// return dense_1(h) # [B, T_enc, d_model] (no relu) -// -// In ggml (channel-fast): -// mel_in ne=[T_mel, n_mels, 1, B] -// dense_0 mul_mat in [n_mels, T_mel, B] orientation -> [d_model, T_mel, B] -// ... but mel_in is [T_mel, n_mels, 1, B]; we permute to -// [n_mels, T_mel, 1, B] first so mul_mat sees n_mels at ne[0]. -// relu, then conv_0 takes [d_model, T_mel, B] as data with kernel -// ne=[k=5, d_model, d_model]; conv_1d_f32 produces [T1, d_model, B] — -// the same transpose pattern gigaam uses. -// Pre-relu, transpose back to [d_model, T_mel, B] for the same-conv shape. -// -// We avoid back-and-forth permutes by keeping the running tensor in -// "channel-fast" [C, T, B] layout for the dense layers (mul_mat) and -// "time-fast" [T, C, B] layout for the conv layers (conv_1d_f32). One -// permute on each boundary is enough. ggml_tensor * build_subsampling(ggml_context * ctx, const MedAsrSubsampling & ss, ggml_tensor * mel_in, @@ -414,7 +365,7 @@ ggml_tensor * build_block(ggml_context * ctx, const float conv_w0 = hp.enc_conv_resid_w0; // 2.0 const float conv_w1 = hp.enc_conv_resid_w1; // 1.0 - // FF1: x = w0*x + w1*FF1(LN(x)). + // Macaron FF1. { ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_ff1_w, ln_eps); y = lasr_feed_forward(ctx, y, b.ff1_up_w, b.ff1_down_w); @@ -422,7 +373,7 @@ ggml_tensor * build_block(ggml_context * ctx, } if (obs != nullptr) obs->post_ff1 = x; - // Self-attention: x = x + Attn(LN(x)). + // Self-attention. { ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_attn_w, ln_eps); y = build_rope_attn(ctx, y, positions, b, d_model, n_head, @@ -431,7 +382,7 @@ ggml_tensor * build_block(ggml_context * ctx, } if (obs != nullptr) obs->post_attn = x; - // Conv: x = w0*x + w1*Conv(LN(x)). + // Conv module. { ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_conv_w, ln_eps); y = build_conv_module(ctx, y, b, bn_scale, bn_bias, @@ -440,7 +391,7 @@ ggml_tensor * build_block(ggml_context * ctx, } if (obs != nullptr) obs->post_conv = x; - // FF2: x = w0*x + w1*FF2(LN(x)). + // Macaron FF2. { ggml_tensor * y = lasr_layer_norm(ctx, x, b.norm_ff2_w, ln_eps); y = lasr_feed_forward(ctx, y, b.ff2_up_w, b.ff2_down_w); diff --git a/src/arch/medasr/model.cpp b/src/arch/medasr/model.cpp index ebd1af85..2ffc3c8e 100644 --- a/src/arch/medasr/model.cpp +++ b/src/arch/medasr/model.cpp @@ -11,10 +11,6 @@ // per-block fused-BN scale/bias inputs, computes the graph, reads // back the CTC logits, performs greedy collapse (drop blanks, collapse // repeats), and populates the result snapshot. -// -// run_batch is left nullptr — the central dispatcher falls back to -// calling run() once per utterance. A run_batch fast path can be wired -// later (Stage 4 Step 9) without changing this file's contract. #include "medasr.h" @@ -303,9 +299,8 @@ transcribe_status init_context(transcribe_model * model, return TRANSCRIBE_OK; } -// Load reference mel sidecar from /mel.in.f32 (validate.py emits -// JSON; the .f32 sidecar is written next to it for fast bytes-only -// reload). For the encoder bring-up path we accept either filename. +// Load reference mel sidecar from /mel.in.f32 (or +// frontend.mel.out.f32). transcribe_status load_ref_mel(const std::string & dir, int n_mels, int & n_mel_frames, @@ -358,12 +353,9 @@ transcribe_status load_ref_mel(const std::string & dir, // Greedy CTC decode: argmax per frame, collapse repeats, drop blanks // AND skip the LASR special tokens (=1, =2, =3) — matching // the reference's `processor.batch_decode(skip_special_tokens=True)`. -// Without this, an utterance whose final CTC frame argmaxes to id 2 -// (``) leaks the literal string `` into the transcript, which -// scores as a +1 insertion against the reference text and lifts the -// dataset WER by ~0.2pp on test-clean. `` (id 0) is already -// dropped by the standard blank-skip; `` and `` are explicit -// skips here to mirror the HF tokenizer's behaviour. +// Without this, a final frame that argmaxes to id 2 (``) leaks the +// literal `` into the transcript, scoring as a +1 insertion (~0.2pp +// WER on test-clean). `` (id 0) is the standard blank-skip. void decode_ctc_greedy(const float * logits, int vocab, int T_enc, @@ -384,10 +376,7 @@ void decode_ctc_greedy(const float * logits, if (best == blank_id) { prev = -1; continue; } if (best == prev) continue; prev = best; - // skip_special_tokens=True equivalent: drop ids 1 (), 2 (), - // and 3 (). The blank id is checked above; this conditional - // is unreachable when blank_id is itself in {1,2,3}, but the - // explicit check keeps the intent local. + // skip_special_tokens=True: drop ids 1 (), 2 (), 3 (). if (best == 1 || best == 2 || best == 3) continue; tokens.push_back(best); frames.push_back(t); diff --git a/src/arch/medasr/weights.cpp b/src/arch/medasr/weights.cpp index 99a0b785..4eaa9a26 100644 --- a/src/arch/medasr/weights.cpp +++ b/src/arch/medasr/weights.cpp @@ -1,11 +1,9 @@ // arch/medasr/weights.cpp - read_medasr_hparams + build_medasr_weights. // -// Pattern follows arch/gigaam/weights.cpp: read every required KV from -// the GGUF, then validate the tensor catalog by name + shape using -// transcribe::weights::find_tensor. Adds fuse_batch_norm() which the -// load() handler calls after data streaming completes — the conv-module -// BatchNorm in each block is fused on the host to a per-channel scale + -// bias pair, and the raw BN tensors stop being consulted after that. +// Reads every required KV, validates the tensor catalog by name + shape, +// then fuse_batch_norm() (called by load() after data streaming) folds +// each block's conv-module BatchNorm on the host to a per-channel scale + +// bias pair, after which the raw BN tensors stop being consulted. #include "weights.h" diff --git a/src/arch/medasr/weights.h b/src/arch/medasr/weights.h index 3bf614de..608b7178 100644 --- a/src/arch/medasr/weights.h +++ b/src/arch/medasr/weights.h @@ -1,21 +1,11 @@ // arch/medasr/weights.h - MedASR (Google LASR-CTC) tensor catalog and // per-instance weight slots. // -// INTERNAL to src/arch/medasr/. Patterned after arch/gigaam/weights.h -// (encoder-CTC Conformer + rotary). MedASR differs in: -// - Subsampling is LINEAR -> RELU -> CONV1D(s=2,p=0) -> RELU -> -// CONV1D(s=2,p=0) -> RELU -> LINEAR (NOT gigaam's 2-conv stack). -// Effective downsampling = 4x. -// - Every LayerNorm in the encoder is bias=false; norm_b slots are -// left nullptr. -// - Self-attn / FF linears are bias=false. -// - Conv module uses BatchNorm (running_mean / running_var fused at -// load) — NOT LayerNorm. -// - Macaron FF residual scalars are [1.5, 0.5] (NOT 0.5 macaron). -// Conv-module residual scalars are [2.0, 1.0]. Per-block scalars -// are baked into the encoder graph builder. -// - CTC head Conv1d(512, 512, k=1) with bias (NOT no-bias Linear). -// CTC blank id = 0 (NOT vocab_size - 1). +// INTERNAL to src/arch/medasr/. Encoder-CTC Conformer + rotary, with a +// 4x subsampling stem (LINEAR -> CONV1D(s=2) -> CONV1D(s=2) -> LINEAR), +// bias=false LayerNorms / linears, BatchNorm conv module (fused at load), +// macaron FF residual scalars [1.5, 0.5] and conv [2.0, 1.0], and a +// Conv1d(k=1) CTC head with bias (blank id = 0). #pragma once diff --git a/src/arch/moonshine/decoder.cpp b/src/arch/moonshine/decoder.cpp index df4e7f86..1b366b71 100644 --- a/src/arch/moonshine/decoder.cpp +++ b/src/arch/moonshine/decoder.cpp @@ -61,10 +61,8 @@ ggml_tensor * unpad_head_dim(ggml_context * ctx, ggml_tensor * t, t->nb[1], t->nb[2], 0); } -// Moonshine uses GPT-J / interleaved RoPE (pairs (0,1), (2,3), ...); -// that is GGML_ROPE_TYPE_NORMAL, not NEOX. See encoder.cpp's -// apply_partial_rope for the full reference (HF rotate_half uses -// `x[..., 0::2]` / `x[..., 1::2]`). +// GPT-J / interleaved RoPE (rotate_half slices x[..., 0::2] / x[..., 1::2]) +// = GGML_ROPE_TYPE_NORMAL, not NEOX. See encoder.cpp apply_partial_rope. ggml_tensor * apply_partial_rope(ggml_context * ctx, ggml_tensor * x, ggml_tensor * positions, @@ -84,19 +82,10 @@ ggml_tensor * apply_partial_rope(ggml_context * ctx, /*beta_slow=*/1.0f); } -// SwiGLU MLP for the decoder: fc1 hidden->2·intermediate, split into -// [x_proj, gate], y = silu(gate) * x_proj, fc2 intermediate->hidden. -// -// HF reference (modeling_moonshine.py:81-88): -// hidden_states = self.fc1(hidden_states) -// hidden_states, gate = hidden_states.chunk(2, dim=-1) -// hidden_states = self.activation_fn(gate) * hidden_states -// hidden_states = self.fc2(hidden_states) -// -// `chunk(2, dim=-1)` splits the last dim of [..., 2·inter] into -// [first inter, second inter]. The first chunk is x_proj, the second is -// gate. In ggml ne the last dim is ne[0] (innermost), so the split is -// the first ne[0]/2 elements (x_proj) and the second ne[0]/2 (gate). +// SwiGLU MLP: fc1 hidden->2·intermediate, split into [x_proj, gate], +// y = silu(gate) * x_proj, fc2 intermediate->hidden. HF chunk(2, dim=-1) +// (modeling_moonshine.py) is the innermost dim = ggml ne[0]: first +// ne[0]/2 is x_proj, second is gate. ggml_tensor * ffn_decoder_swiglu(ggml_context * ctx, ggml_tensor * x, ggml_tensor * fc1_w, ggml_tensor * fc1_b, @@ -159,11 +148,10 @@ ggml_tensor * mha_self_cached(ggml_context * ctx, const int head_dim_rot = hp.dec_head_dim_rot(); const int pad = head_dim_pad - head_dim; const int n_ctx = kv_cache.n_ctx; - // Unpadded scale: HF uses self.head_dim ** -0.5; padded slots are - // zero-extension and don't change the dot product magnitude. + // Unpadded scale: HF uses head_dim**-0.5; padded slots are + // zero-extension and don't change the dot product. const float scale = 1.0f / std::sqrt(static_cast(head_dim)); - // Q, K, V projections (bias-less). ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x); ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x); ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x); @@ -178,11 +166,9 @@ ggml_tensor * mha_self_cached(ggml_context * ctx, ggml_tensor * Q_rope = apply_partial_rope(ctx, Qcur, pos_ids, hp, head_dim_rot); ggml_tensor * K_rope = apply_partial_rope(ctx, Kcur, pos_ids, hp, head_dim_rot); - // Q needs [head_dim, n_tokens, n_heads, 1] for the attention path. - // K and V do NOT — they're cached in [head_dim, n_heads, n_tokens] - // memory order, which is exactly what K_rope / Vcur already have - // after reshape_4d + (optional) RoPE. Skip the round-trip permute - // + cont that the previous code did before writing to the cache. + // Q needs [head_dim, n_tokens, n_heads, 1] for attention. K/V are + // already in [head_dim, n_heads, n_tokens] cache memory order after + // reshape + (optional) RoPE, so they're written straight to the cache. ggml_tensor * Q_unpad = ggml_cont(ctx, ggml_permute(ctx, Q_rope, 0, 2, 1, 3)); // Write K_rope / Vcur into the cache. The cache stores @@ -260,7 +246,6 @@ ggml_tensor * mha_self_cached(ggml_context * ctx, o = ggml_cont(ctx, o); o = ggml_reshape_2d(ctx, o, d_model, n_tokens); - // Output projection (bias-less). o = ggml_mul_mat(ctx, out_w, o); return o; } @@ -298,7 +283,6 @@ ggml_tensor * mha_self_step(ggml_context * ctx, const int n_ctx = kv_cache.n_ctx; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); - // Q, K, V projections (bias-less). ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x); ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x); ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x); @@ -467,10 +451,7 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, } // namespace -// --------------------------------------------------------------------------- -// Cross-KV precompute graph -// --------------------------------------------------------------------------- - +// Cross-KV precompute graph. DecoderBuild build_cross_kv_graph(ggml_context * ctx, const MoonshineWeights & w, const MoonshineHParams & hp, @@ -531,10 +512,7 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- -// KV-cached decoder graph (prompt + step) -// --------------------------------------------------------------------------- - +// KV-cached decoder graph (prompt + step). DecoderBuild build_decoder_graph_kv(ggml_context * ctx, const MoonshineWeights & w, const MoonshineHParams & hp, @@ -597,14 +575,10 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, } // Moonshine has no additive positional embedding; embed_sum is - // identical to token_emb. We still emit it under the "embed_sum" - // name so the validate.py harness sees the same tensor stream as - // whisper / cohere. + // identical to token_emb, still emitted under "embed_sum" for dump + // parity with whisper / cohere. ggml_tensor * x = tok_emb; { - // ggml_set_name lives on the same tensor; we re-use named() to - // tag it. Note: the tensor identity is the same so dumping - // either name reads the same memory. ggml_tensor * embed_sum = ggml_scale(ctx, tok_emb, 1.0f); named(embed_sum, "dec.embed_sum"); if (n_past == 0) { @@ -626,7 +600,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, for (int i = 0; i < n_blocks; ++i) { const auto & b = w.dec_blocks[i]; - // Self-attn (pre-LN, partial RoPE, causal via mask when n_tokens>1). { ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, /*beta=*/nullptr); y = mha_self_cached( @@ -636,7 +609,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, i, n_past, n_tokens, n_kv, use_flash); x = ggml_add(ctx, x, y); } - // Cross-attn (pre-LN, no RoPE). { ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, /*beta=*/nullptr); y = mha_cross_cached( @@ -645,7 +617,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, hp, n_heads, d_model, i, T_enc, use_flash); x = ggml_add(ctx, x, y); } - // SwiGLU MLP (pre-LN). { ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, /*beta=*/nullptr); y = ffn_decoder_swiglu(ctx, y, @@ -672,9 +643,8 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.out_before_head = x; } - // Tied lm_head: proj_out reuses dec.token_embd.weight. mul_mat - // expects W ne=[d_model, vocab], x ne=[d_model, n_tokens]; result - // ne=[vocab, n_tokens]. + // Tied lm_head (reuses dec.token_embd.weight): W ne=[d_model, vocab], + // x ne=[d_model, n_tokens] -> logits ne=[vocab, n_tokens]. ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x); named(logits_raw, "dec.logits_raw"); if (n_past == 0) { @@ -695,10 +665,9 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.out = logits; ggml_set_output(db.out); - // When skipping log_softmax (every step pass — every call after the - // dump-emitting prompt pass), append argmax over the last position so - // we only download a single int32 instead of the full vocab logits. - // Mirrors cohere/decoder.cpp:756-771. + // When skipping log_softmax (every step pass after the dump-emitting + // prompt pass), append argmax over the last position so we download a + // single int32 instead of the full vocab logits. if (skip_log_softmax) { ggml_tensor * last_logits = logits; if (n_tokens > 1) { @@ -722,10 +691,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- // Static-topology single-token step graph (GPU dispatch path). -// --------------------------------------------------------------------------- - StepBuild build_step_graph(ggml_context * ctx, const MoonshineWeights & w, const MoonshineHParams & hp, @@ -782,7 +748,6 @@ StepBuild build_step_graph(ggml_context * ctx, for (int i = 0; i < n_blocks; ++i) { const auto & b = w.dec_blocks[i]; - // Self-attn (pre-LN, partial RoPE, set_rows-based KV write). { ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, /*beta=*/nullptr); y = mha_self_step( @@ -792,7 +757,6 @@ StepBuild build_step_graph(ggml_context * ctx, i, max_n_kv, use_flash); x = ggml_add(ctx, x, y); } - // Cross-attn (pre-LN, no RoPE, reads pre-populated cache). { ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, /*beta=*/nullptr); y = mha_cross_cached( @@ -801,7 +765,6 @@ StepBuild build_step_graph(ggml_context * ctx, hp, n_heads, d_model, i, T_enc, use_flash); x = ggml_add(ctx, x, y); } - // SwiGLU MLP (pre-LN). { ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, /*beta=*/nullptr); y = ffn_decoder_swiglu(ctx, y, @@ -825,10 +788,7 @@ StepBuild build_step_graph(ggml_context * ctx, return sb; } -// =========================================================================== -// Offline batched decode (B utterances) -// =========================================================================== - +// Offline batched decode (B utterances). namespace { // Slice off head-dim padding on a 4D tensor [head_dim_pad, a, b, c] → keep B. diff --git a/src/arch/moonshine/decoder.h b/src/arch/moonshine/decoder.h index 5cd1c5bc..07ace35f 100644 --- a/src/arch/moonshine/decoder.h +++ b/src/arch/moonshine/decoder.h @@ -11,8 +11,7 @@ // Handles both the prompt pass (n_tokens=1, n_past=0 — moonshine's // prompt is a single bos token) and steady-state step passes // (n_tokens=1, n_past=current). All dec.* dump points are wired -// on the prompt pass; the validate.py harness compares them -// against the reference dumps. +// on the prompt pass. // // All attention projections (q/k/v/o) are bias-less (attention_bias=False). // Self-attn applies partial RoPE 0.9 to q/k; cross-attn does NOT rotate. @@ -116,10 +115,8 @@ StepBuild build_step_graph(ggml_context * compute_ctx, int T_enc, bool use_flash = true); -// --------------------------------------------------------------------------- // Offline batched decode (B utterances). Mirrors src/arch/cohere + canary, // with moonshine's partial RoPE on self-attn and head-dim padding. -// --------------------------------------------------------------------------- // Batched cross-attention K/V from packed encoder outputs [d_model, T_enc_max, B]. DecoderBuild build_cross_kv_graph_batched(ggml_context * ctx, diff --git a/src/arch/moonshine/encoder.cpp b/src/arch/moonshine/encoder.cpp index 7866c0b4..d7996152 100644 --- a/src/arch/moonshine/encoder.cpp +++ b/src/arch/moonshine/encoder.cpp @@ -1,42 +1,14 @@ // arch/moonshine/encoder.cpp - Moonshine encoder graph builder. // -// Forward shape: +// 3-conv stem (conv0 k=127 s=64 + tanh + GroupNorm; conv1 k=7 s=3 + gelu; +// conv2 k=3 s=2 + gelu) on raw 16 kHz PCM, then n_layers pre-LN blocks +// (bidirectional partial-RoPE MHSA + GELU MLP) and a final bias-less LN. // -// audio [T_samples] (raw 16 kHz PCM in [-1, 1]) -// -> reshape [T_samples, in=1] for ggml_conv_1d -// -> conv0 (k=127, s=64, no bias) [T1, d_model] -// -> tanh -// -> transpose to [d_model, T1] (channel-innermost, -// matching reference dump -// layout enc.conv1.out) -// -> dump enc.conv1.out -// -> bias-less LayerNorm (= GroupNorm num_groups=1, ε=1e-5) -// with affine (* gn_w + gn_b) -// -> dump enc.groupnorm.out -// -> transpose back to [T1, d_model] for conv1 -// -> conv1 (k=7, s=3) + bias [T2, 2·d_model] -// -> gelu -// -> transpose to [2·d_model, T2] -// -> dump enc.conv2.out -// -> transpose back to [T2, 2·d_model] for conv2 -// -> conv2 (k=3, s=2) + bias [T_enc, d_model] -// -> gelu -// -> transpose to [d_model, T_enc] -// -> dump enc.conv3.out -// -> n_layers x pre-LN transformer block (partial-RoPE MHSA + GELU MLP) -// -> final layer_norm (no bias) -// -> dump enc.final +// HF uses exact-erf gelu (ggml_gelu_erf); the tanh approx drifts ~1e-4/elem. // -// Conv stem note: HF uses `nn.functional.gelu(...)` (exact-erf form) for -// every gelu call in moonshine. ggml_gelu_erf matches; ggml_gelu is the -// tanh approximation and drifts ~1e-4 per element. -// -// Encoder block: pre-LN MHSA (no causal mask — bidirectional, no mask -// at all) + pre-LN GELU MLP. Self-attn applies partial RoPE to the -// leading `head_dim_rot` dims of q/k. q/k/v are pre-padded to -// `head_dim_padded = round_up(head_dim, pad_head_dim_multiple)` before -// the matmul; the trailing `head_dim_padding` rows of attn_output are -// sliced off before o_proj (matching HF MoonshineAttention). +// q/k/v are pre-padded to head_dim_padded = round_up(head_dim, +// pad_head_dim_multiple); the padding rows of attn_output are sliced off +// before o_proj (matching HF MoonshineAttention). #include "encoder.h" @@ -70,15 +42,10 @@ ggml_tensor * add_conv1d_bias(ggml_context * ctx, return ggml_add(ctx, conv_out, bias_4d); } -// Apply partial RoPE to `x` (shape [head_dim, n_heads, T, 1] in ggml -// ne — head_dim at ne[0], T at ne[2]). RoPE rotates the leading -// `head_dim_rot` of dim-0; the trailing dims pass through unchanged. -// -// Moonshine uses GPT-J / **interleaved** RoPE: rotation pairs are -// (0, 1), (2, 3), ... — see `rotate_half` in HF modeling_moonshine.py -// (`x[..., 0::2]` / `x[..., 1::2]`). That is GGML_ROPE_TYPE_NORMAL -// (mode=0), NOT GGML_ROPE_TYPE_NEOX (mode=2, which uses split-halves -// pairing (0, D/2), (1, D/2+1), ...). +// Partial RoPE on the leading `head_dim_rot` of dim-0 (x ne = +// [head_dim, n_heads, T, 1]). Moonshine uses GPT-J / interleaved RoPE +// (rotate_half slices x[..., 0::2] / x[..., 1::2]) = GGML_ROPE_TYPE_NORMAL, +// NOT NEOX. ggml_tensor * apply_partial_rope(ggml_context * ctx, ggml_tensor * x, ggml_tensor * positions, @@ -119,14 +86,11 @@ ggml_tensor * mha_encoder(ggml_context * ctx, const int head_dim_pad = hp.enc_head_dim_padded(); const int head_dim_rot = hp.enc_head_dim_rot(); const int pad = head_dim_pad - head_dim; - // HF MoonshineAttention uses unpadded head_dim for the scale - // (`self.scaling = self.head_dim**-0.5`). The padded slots are - // zero-extension and contribute zero to the dot product, so the - // unpadded scale is the numerically correct match. + // Unpadded scale: HF uses head_dim**-0.5; padded slots are + // zero-extension and don't change the dot product. const float scale = 1.0f / std::sqrt(static_cast(head_dim)); const int64_t T = x->ne[1]; - // Q, K, V projections (bias-less). ggml_tensor * q = ggml_mul_mat(ctx, q_w, x); // [d_model, T] ggml_tensor * k = ggml_mul_mat(ctx, k_w, x); ggml_tensor * v = ggml_mul_mat(ctx, v_w, x); @@ -158,7 +122,6 @@ ggml_tensor * mha_encoder(ggml_context * ctx, k = to_attn_layout(k); v = to_attn_layout(v); - // Attention. ggml_tensor * o; if (use_flash) { o = ggml_flash_attn_ext(ctx, q, k, v, /*mask=*/nullptr, @@ -177,10 +140,8 @@ ggml_tensor * mha_encoder(ggml_context * ctx, // o shape: [head_dim_pad, T, n_heads, 1] } - // Slice off the head_dim padding before merging heads (HF order: - // slice in [..., :head_dim] BEFORE the n_heads*head_dim merge so the - // padding regions of every head are removed). In ggml ne terms the - // op is ggml_view_3d on dim-0 with new ne[0]=head_dim, then cont. + // Slice off head_dim padding before merging heads (HF slices + // [..., :head_dim] before the n_heads*head_dim merge). if (pad > 0) { o = ggml_view_3d(ctx, o, head_dim, T, n_heads, @@ -193,7 +154,6 @@ ggml_tensor * mha_encoder(ggml_context * ctx, o = ggml_cont(ctx, o); o = ggml_reshape_2d(ctx, o, d_model, T); - // Output projection (bias-less). o = ggml_mul_mat(ctx, out_w, o); return o; } @@ -221,7 +181,6 @@ ggml_tensor * build_block(ggml_context * ctx, int d_model, bool use_flash) { - // Self-attn sublayer (bias-less LN, bias-less projections). { ggml_tensor * y = layer_norm(ctx, x, b.norm_attn_w, /*beta=*/nullptr); y = mha_encoder(ctx, y, pos_ids, @@ -229,7 +188,6 @@ ggml_tensor * build_block(ggml_context * ctx, hp, n_heads, d_model, use_flash); x = ggml_add(ctx, x, y); } - // FFN sublayer. { ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, /*beta=*/nullptr); y = ffn_encoder(ctx, y, @@ -306,16 +264,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_set_input(eb.pos_ids_in); // ---- conv stem ---- - // ggml_conv_1d wants [T, in_channels]. PCM is 1 channel: reshape - // [n_samples] -> [n_samples, 1]. + // ggml_conv_1d / conv_1d_f32 want input ne=[T, in_channels] (T + // innermost). PCM is 1 channel: reshape [n_samples] -> [n_samples, 1]. ggml_tensor * x = ggml_reshape_2d(ctx, eb.audio_in, n_samples, 1); - // Now x ne=[n_samples, 1]. ggml_conv_1d expects ne[0]=T innermost, - // ne[1]=in_channels. We need [T, in] with T innermost — actually - // checking conformer's conv_1d_f32 helper signature: it takes input - // ne=[T, Cin] (T innermost). ggml_reshape_2d sets ne[0]=n_samples, - // ne[1]=1 — that's [T, 1] with T innermost. ✓ - // conv0: k=127, s=64, no padding, no bias. x = conf::conv_1d_f32(ctx, w.enc_stem.conv0_w, x, /*s=*/hp.conv_strides[0], /*p=*/0, @@ -330,18 +282,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.conv1_out = x; transcribe::debug::mark_tensor_for_dump(x); - // GroupNorm with num_groups=1 over [B, C, T]: PyTorch normalizes - // jointly over (C, T) per batch (NOT just over C, like LayerNorm). - // ggml_group_norm with n_groups=1 matches: it pools all elements - // outside the leading group axis into a single moment computation. - // Affine (* gn_w + gn_b) is applied separately — ggml_group_norm - // handles only the standardization. - // - // Layout note: ggml_group_norm reads ne[2] as the channel axis - // (matching stable-diffusion's [N, H, W, C] convention with C at - // ne[0]). We currently have x in ne=[C, T1] = ne=[288, 2749]; need - // to reshape to ne=[C, T1, 1, 1] so the group axis (ne[2]) is the - // expected one. (ggml internally splits ne[2] into n_groups.) + // GroupNorm num_groups=1: PyTorch normalizes jointly over (C, T) per + // batch (not just over C). ggml_group_norm matches (standardization + // only; affine * gn_w + gn_b applied separately below), but reads + // ne[2] as the channel axis, so reshape [C, T1] -> [C, T1, 1, 1]. { const int64_t C = w.enc_stem.gn_w->ne[0]; const int64_t T1 = x->ne[1]; @@ -359,7 +303,6 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Transpose back to [T1, d_model] for conv1. x = ggml_cont(ctx, ggml_transpose(ctx, x)); - // conv1: k=7, s=3, with bias. x = conf::conv_1d_f32(ctx, w.enc_stem.conv1_w, x, /*s=*/hp.conv_strides[1], /*p=*/0, @@ -375,7 +318,6 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Back to [T2, 2·d_model] for conv2. x = ggml_cont(ctx, ggml_transpose(ctx, x)); - // conv2: k=3, s=2, with bias. x = conf::conv_1d_f32(ctx, w.enc_stem.conv2_w, x, /*s=*/hp.conv_strides[2], /*p=*/0, diff --git a/src/arch/moonshine/encoder.h b/src/arch/moonshine/encoder.h index 51a737d9..4f52141e 100644 --- a/src/arch/moonshine/encoder.h +++ b/src/arch/moonshine/encoder.h @@ -1,21 +1,14 @@ // arch/moonshine/encoder.h - Moonshine encoder graph builder. // -// The encoder is a 3-layer Conv1d stem on raw 16 kHz PCM (kernels -// 127/7/3, strides 64/3/2; conv0 has no bias, conv1/conv2 do; conv0 -// is followed by tanh + GroupNorm(num_groups=1, eps=1e-5); conv1/conv2 -// each followed by GELU). After the stem, the encoder runs `n_layers` -// pre-LN transformer blocks with partial-RoPE self-attention (no -// causal mask — bidirectional) and GELU MLP. A final bias-less -// LayerNorm closes the encoder. +// 3-conv stem on raw 16 kHz PCM (see encoder.cpp / weights.h) + n_layers +// pre-LN bidirectional partial-RoPE transformer blocks + final bias-less LN. // -// Dump points (must match `build/validate/moonshine//dump_coverage.json`): +// Dump points: // enc.audio.in raw PCM input // enc.conv1.out after tanh(conv0) (BEFORE GroupNorm) // enc.groupnorm.out after GroupNorm // enc.conv2.out after gelu(conv1) // enc.conv3.out after gelu(conv2) + permute (this is the block input) -// enc.rope.cos partial-RoPE cosine table [T_enc, head_dim_rot] -// enc.rope.sin partial-RoPE sine table [T_enc, head_dim_rot] // enc.block.{i}.out per-block residual stream // enc.final after final LN diff --git a/src/arch/moonshine/model.cpp b/src/arch/moonshine/model.cpp index 34fe9f57..8accd8dc 100644 --- a/src/arch/moonshine/model.cpp +++ b/src/arch/moonshine/model.cpp @@ -75,29 +75,16 @@ MoonshineModel::~MoonshineModel() { plan.primary_kind = transcribe::BackendKind::Unknown; } -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Moonshine is the honest -// edge case in the soft-window bucket: its only wall is on *output*, not -// input. The conv-stem encoder is effectively unbounded (it processes any -// PCM length), so a long clip is never rejected. Instead the greedy decode -// loop stops when it hits the decoder's position cap -// (dec_max_position_embeddings = 194 for moonshine) before emitting EOS, and -// the transcript silently truncates. We do not gate on audio length — a dense -// short clip can hit the cap too — so there is no upfront INPUT_TOO_LONG -// rejection here. Instead, when a decode loop exits on the cap rather than on -// EOS, we set transcribe_was_truncated() and emit a WARN (same convention as -// qwen3_asr's generation-budget guard). -// --------------------------------------------------------------------------- - -// Advisory transcribe_capabilities::max_audio_ms. Because the limit is -// output-bound and content-dependent (it counts decode tokens, not audio -// frames), there is no exact audio-length ceiling. We report a conservative -// advisory: the audio duration the output budget (dec_max_position_embeddings -// tokens) can realistically cover at a typical speaking rate of ~4 output -// tokens/sec. For moonshine (194 tokens) that is ~48 s — moonshine is intended -// for short utterances. Returns 0 ("unknown / unbounded") if the cap is -// missing. Crossing this advisory does not reject the clip; it only makes a -// mid-decode truncation (transcribe_was_truncated) more likely. +// Input-length contract (see docs/input-limits.md): the wall is on *output*, +// not input. The conv-stem encoder takes any PCM length, so a clip is never +// rejected; the greedy decode loop instead stops at the decoder position cap +// (dec_max_position_embeddings = 194) before EOS and the transcript silently +// truncates. On a cap-exit we set transcribe_was_truncated() and WARN (same as +// qwen3_asr). max_audio_ms below is a conservative advisory only. + +// Advisory transcribe_capabilities::max_audio_ms: the audio the output budget +// (dec_max_position_embeddings tokens) covers at ~4 output tokens/sec (~48 s +// for 194 tokens). 0 means unknown/unbounded. Advisory only — does not reject. constexpr int k_tokens_per_sec = 4; // rough speech-rate estimate; advisory only int64_t moonshine_max_audio_ms(const MoonshineHParams & hp) { if (hp.dec_max_position_embeddings <= 0) { @@ -107,10 +94,7 @@ int64_t moonshine_max_audio_ms(const MoonshineHParams & hp) { k_tokens_per_sec; } -// --------------------------------------------------------------------------- -// KV cache initialization -// --------------------------------------------------------------------------- - +// KV cache initialization. bool kv_cache_init(MoonshineKvCache & cache, ggml_backend_t backend, int n_ctx, @@ -248,12 +232,11 @@ transcribe_status load( if (auto st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) return st; if (auto st = read_moonshine_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) return st; - // Publish the advisory input-length window now that the decoder position - // cap is known. apply_family_invariants() ran above (before hparams), so - // this is the first point the cap is available. See docs/input-limits.md. + // 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_max_audio_ms(m->hparams); - // Stage 2: reopen GGUF with no_alloc to wire weight slots. + // Reopen GGUF with no_alloc to wire weight slots. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -427,7 +410,6 @@ transcribe_status run( return TRANSCRIBE_ERR_OOM; } - // Upload PCM and encoder pos_ids. ggml_backend_tensor_set(eb.audio_in, pcm, 0, static_cast(n_samples) * sizeof(float)); { @@ -444,7 +426,6 @@ transcribe_status run( return TRANSCRIBE_ERR_GGUF; } - // Dump encoder intermediates. auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { if (t != nullptr) transcribe::debug::dump_tensor(name, t, stage); }; @@ -479,8 +460,7 @@ transcribe_status run( const int n_ctx = hp.dec_max_position_embeddings > 0 ? hp.dec_max_position_embeddings : 512; ggml_type cache_type = resolved_kv; - // For F32 reference dtype, default the cache to F32 too — - // matching the reference numerical regime in Stage 4. + // 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, @@ -580,8 +560,7 @@ transcribe_status run( 0, token_ids.size() * sizeof(int32_t)); ggml_backend_tensor_set(db.pos_ids_in, pos_ids.data(), 0, pos_ids.size() * sizeof(int32_t)); - // Causal mask only when n_tokens > 1 (not exercised on this - // family today — every step is single-token). + // Causal mask only when n_tokens > 1 (every step is single-token). if (n_tokens > 1) { ggml_tensor * m = find_tensor_by_name(cc->compute_ctx, "dec.causal_mask"); if (m != nullptr) { @@ -795,14 +774,9 @@ transcribe_status run( } } - // Both greedy paths above (GPU static-graph loop and CPU/debug dynamic - // loop) exit either because next_token == eos (complete) or because - // n_past reached the position cap (max_pos) — or, on the GPU path, the - // KV-width break, which only fires once n_past has reached max_pos. When - // the last token was not eos the transcript hit the output cap before - // end-of-stream and is incomplete: flag it via transcribe_was_truncated() - // and emit one WARN, mirroring qwen3_asr. The decode loops themselves are - // unchanged — this only observes their exit. See docs/input-limits.md. + // A non-eos last token means the decode hit the position cap before + // end-of-stream (see the input-length contract above): flag truncation + // and WARN. if (next_token != eos) { cc->was_truncated = true; transcribe::log_msg( @@ -816,7 +790,6 @@ transcribe_status run( cc->t_decode_us = ggml_time_us() - t_decode_start; - // Decode IDs to text. if (!generated_ids.empty()) { std::string full = cm->tok.decode(generated_ids.data(), static_cast(generated_ids.size())); @@ -842,23 +815,16 @@ transcribe_status run( cc->has_result = true; } - // Output truncation is a hard status (see docs/input-limits.md): the - // partial transcript above stays readable, but the caller is told the - // run did not reach end-of-stream. was_truncated was set in the - // OFFLINE single-utterance decode-exit check above. + // Truncation is a hard status; the partial transcript stays readable. return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } -// =========================================================================== // Offline batched decode (transcribe_run_batch). Mirrors src/arch/cohere + -// canary. Encoder serial per utterance (compute-bound; moonshine has no mel -// frontend to parallelize — raw PCM passthrough); decode (self+cross attn, -// partial RoPE) batched. Requires the flash step path, so on Metal — where -// moonshine defaults decoder flash OFF (head_dim_padded unsupported) — this -// falls back to serial; it engages on Vulkan/CUDA. See -// docs/batching-autoregressive-plan.md. -// =========================================================================== +// canary. Encoder serial per utterance (raw PCM, no mel to parallelize); +// decode (self+cross attn, partial RoPE) batched. Requires the flash step +// path, so on Metal (decoder flash OFF — head_dim_padded unsupported) it falls +// back to serial; engages on Vulkan/CUDA. transcribe_status encode_one_to_host( MoonshineSession * cc, MoonshineModel * cm, @@ -954,11 +920,9 @@ transcribe_status run_batch( cm->plan.primary_kind != transcribe::BackendKind::Cpu && cm->plan.primary_kind != transcribe::BackendKind::Accel && cm->plan.primary_kind != transcribe::BackendKind::Unknown; - // The batched path always uses flash, even on Metal where moonshine - // defaults decoder flash OFF for single-shot (base's head_dim_padded=56 - // CPU-spills). At B>1 the flash path still beats serial (measured: tiny - // ~1.7×, base ~1.17×), so we don't gate on cc->decoder_use_flash here — - // the batched step graph below is built with use_flash=true unconditionally. + // The batched path always uses flash (unconditional), even on Metal where + // single-shot defaults decoder flash OFF (base's head_dim_padded=56 + // CPU-spills): at B>1 flash still beats serial (tiny ~1.7×, base ~1.17×). if (n == 1 || !primary_is_gpu || transcribe::debug::enabled()) return run_batch_serial(cc, pcm, n_samples, n, params); @@ -1107,11 +1071,8 @@ transcribe_status run_batch( } const int64_t dec_us = ggml_time_us() - t_dec0; - // Batched truncation flag. The shared step loop marks each valid row that - // stopped on the output cap (max_pos) / context window before end-of-stream - // — its transcript is incomplete. Mirror the serial path: set - // transcribe_was_truncated() and emit one WARN. (The decode is unchanged.) - // See docs/input-limits.md. + // Batched truncation: the shared step loop marks each valid row that hit + // the output cap before end-of-stream. Mirror the serial path (WARN + flag). { int n_truncated = 0; for (int b = 0; b < n; ++b) { @@ -1147,9 +1108,7 @@ transcribe_status run_batch( rs.full_text = full; rs.result_kind = TRANSCRIBE_TIMESTAMPS_NONE; rs.has_result = true; rs.status = TRANSCRIBE_OK; - // Per-utterance truncation parity with the single-shot path: a valid row - // marked truncated by the shared loop reports - // TRANSCRIBE_ERR_OUTPUT_TRUNCATED (partial transcript retained). Only + // Per-utterance truncation parity with the single-shot path. Only // override an otherwise-OK status — never a worse one. if (rs.status == TRANSCRIBE_OK && b < static_cast(truncated.size()) && truncated[b]) { diff --git a/src/arch/moonshine/moonshine.h b/src/arch/moonshine/moonshine.h index 50bdebf2..8a7d85a6 100644 --- a/src/arch/moonshine/moonshine.h +++ b/src/arch/moonshine/moonshine.h @@ -137,23 +137,12 @@ struct MoonshineSession final : public transcribe_session { // Flash-attention defaults — finalized per-backend in init_context. - // - // Encoder: ON everywhere. The encoder runs as one big graph with - // nq=T_enc, and both shipped variants benefit from FA there (tiny - // especially — ~2x encoder speedup on long audio). - // - // Decoder: ON everywhere except Metal. Vulkan/RADV measures ~15% - // faster on decode_ms with FA on (jfk q8_0: tiny 112→98 ms, - // base 165→143 ms). Metal is the outlier: moonshine-base's - // head_dim_padded=56 is NOT in Metal's supported FA head-size set - // ({32,40,48,64,72,80,96,...}), so the FA op spills per-step to CPU - // and is ~2x SLOWER than the explicit kq->softmax->v·kq path; tiny's - // head_dim_padded=40 IS supported on Metal but only via the tile - // kernel (vec needs head_dim%32==0), roughly a wash. Net: keep - // decoder FA off on Metal until either (a) we re-pad head_dim to - // 64 to enable the FA vec kernel, or (b) ggml Metal grows support - // for these head sizes. TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH - // still apply on top. + // Encoder: ON everywhere (~2x on long audio). Decoder: ON except Metal + // (Vulkan ~15% faster: jfk q8_0 tiny 112→98, base 165→143 ms). Metal + // is the outlier: base's head_dim_padded=56 is not in Metal's FA + // head-size set, so FA spills per-step to CPU and is ~2x slower than the + // explicit kq->softmax->v·kq path (tiny's 40 is supported only via the + // tile kernel, a wash). TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH apply on top. bool encoder_use_flash = true; bool decoder_use_flash = true; diff --git a/src/arch/moonshine_streaming/capabilities.cpp b/src/arch/moonshine_streaming/capabilities.cpp index ca4dd29f..6468df22 100644 --- a/src/arch/moonshine_streaming/capabilities.cpp +++ b/src/arch/moonshine_streaming/capabilities.cpp @@ -19,10 +19,13 @@ void apply_family_invariants(transcribe_model & model) { // text-only segment with zeroed timings (same policy as moonshine). caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; - // Phase 4b-encoder streaming: stream_feed runs the encoder - // incrementally over a sliding window, appending stable frames to - // a per-utterance committed buffer; stream_finalize runs adapter + - // cross_kv + AR decode once over the full committed buffer. + // Streaming: stream_feed runs the encoder, adapter, and cross-KV + // projection incrementally over a sliding window, appending each stable + // slice to per-utterance host K/V buffers, and can run a throttled + // partial AR decode over the growing cross-KV for a live transcript. + // stream_finalize only tops up the trailing tail and runs a final AR + // decode when frames advanced past the last partial (otherwise it just + // commits the last partial transcript). caps.supports_streaming = true; // Streaming latency characteristics (≈240 ms cumulative encoder diff --git a/src/arch/moonshine_streaming/decoder.cpp b/src/arch/moonshine_streaming/decoder.cpp index 76025646..8a785ad5 100644 --- a/src/arch/moonshine_streaming/decoder.cpp +++ b/src/arch/moonshine_streaming/decoder.cpp @@ -11,9 +11,7 @@ // not dec.token_embd.weight. // - Decoder LayerNorms are vanilla `nn.LayerNorm(bias=False)` — no // unit_offset folding. -// - RoPE rotation mode: same as moonshine — GPT-J / interleaved -// (rotate_half slices `x[..., 0::2]` / `x[..., 1::2]`) → -// GGML_ROPE_TYPE_NORMAL, NOT NEOX. +// - RoPE: GPT-J / interleaved (= GGML_ROPE_TYPE_NORMAL, not NEOX), as moonshine. #include "decoder.h" @@ -73,14 +71,8 @@ ggml_tensor * apply_partial_rope(ggml_context * ctx, /*beta_slow=*/1.0f); } -// SwiGLU MLP: fc1 outputs `2·ffn_dim`, chunk into [hidden_states, gate], -// out = silu(gate) * hidden_states, fc2. -// -// HF reference: -// hidden_states = self.fc1(hidden_states) -// hidden_states, gate = hidden_states.chunk(2, dim=-1) -// hidden_states = self.activation_fn(gate) * hidden_states -// hidden_states = self.fc2(hidden_states) +// SwiGLU MLP: fc1 -> 2·ffn_dim, chunk(2, dim=-1) into [x_proj, gate] +// (innermost dim = ggml ne[0]), out = silu(gate) * x_proj, fc2. ggml_tensor * ffn_decoder_swiglu(ggml_context * ctx, ggml_tensor * x, ggml_tensor * fc1_w, ggml_tensor * fc1_b, @@ -143,9 +135,8 @@ ggml_tensor * mha_self_cached(ggml_context * ctx, ggml_tensor * K_rope = apply_partial_rope(ctx, Kcur, pos_ids, hp, head_dim_rot); // Q needs [head_dim, n_tokens, n_heads, 1] for attention. K/V are - // already [head_dim, n_heads, n_tokens, 1] contiguous after RoPE - // (or reshape, for V) — same memory order as the cache slot. Skip - // the round-trip permute+cont before writing. + // already in cache memory order [head_dim, n_heads, n_tokens, 1] after + // RoPE (or reshape, for V), so they're written straight to the cache. ggml_tensor * Q_unpad = ggml_cont(ctx, ggml_permute(ctx, Q_rope, 0, 2, 1, 3)); { @@ -288,10 +279,7 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, } // namespace -// --------------------------------------------------------------------------- -// Adapter graph -// --------------------------------------------------------------------------- - +// Adapter graph. AdapterBuild build_adapter_graph(ggml_context * ctx, const MoonshineStreamingWeights & w, const MoonshineStreamingHParams & hp, @@ -351,10 +339,7 @@ AdapterBuild build_adapter_graph(ggml_context * ctx, return ab; } -// --------------------------------------------------------------------------- -// Cross-KV precompute graph -// --------------------------------------------------------------------------- - +// Cross-KV precompute graph. DecoderBuild build_cross_kv_graph(ggml_context * ctx, const MoonshineStreamingWeights & w, const MoonshineStreamingHParams & hp, @@ -412,10 +397,7 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- -// Cross-KV projection (incremental streaming) -// --------------------------------------------------------------------------- - +// Cross-KV projection (incremental streaming). CrossKVProjectionBuild build_cross_kv_projection_graph( ggml_context * ctx, const MoonshineStreamingWeights & w, @@ -472,10 +454,7 @@ CrossKVProjectionBuild build_cross_kv_projection_graph( return pb; } -// --------------------------------------------------------------------------- -// Cross-KV commit (host buffers → persistent cache) -// --------------------------------------------------------------------------- - +// Cross-KV commit (host buffers → persistent cache). CrossKVCommitBuild build_cross_kv_commit_graph( ggml_context * ctx, const MoonshineStreamingHParams & hp, @@ -542,10 +521,7 @@ CrossKVCommitBuild build_cross_kv_commit_graph( return cb; } -// --------------------------------------------------------------------------- -// KV-cached decoder graph (prompt + step) -// --------------------------------------------------------------------------- - +// KV-cached decoder graph (prompt + step). DecoderBuild build_decoder_graph_kv(ggml_context * ctx, const MoonshineStreamingWeights & w, const MoonshineStreamingHParams & hp, @@ -624,7 +600,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, for (int i = 0; i < n_blocks; ++i) { const auto & b = w.dec_blocks[i]; - // Self-attn (pre-LN, partial RoPE, causal). { ggml_tensor * y = layer_norm(ctx, x, b.norm_self_w, /*beta=*/nullptr); y = mha_self_cached( @@ -634,7 +609,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, i, n_past, n_tokens, n_kv, use_flash); x = ggml_add(ctx, x, y); } - // Cross-attn (pre-LN, no RoPE). { ggml_tensor * y = layer_norm(ctx, x, b.norm_cross_w, /*beta=*/nullptr); y = mha_cross_cached( @@ -643,7 +617,6 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, hp, n_heads, d_model, i, T_enc, use_flash); x = ggml_add(ctx, x, y); } - // SwiGLU MLP (pre-LN). { ggml_tensor * y = layer_norm(ctx, x, b.norm_ffn_w, /*beta=*/nullptr); y = ffn_decoder_swiglu(ctx, y, @@ -655,8 +628,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, if (n_past == 0) { // Track every block so model.cpp can index by layer; only mark - // the auto_blocks subset for set_output so the scheduler keeps - // matching tensors as the reference dumper. + // the auto_blocks subset for dump (see dump_block_index). db.dumps.block_outs.push_back(x); if (dump_block_index(i, n_blocks)) { char bname[64]; @@ -674,8 +646,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.out_before_head = x; } - // Untied lm_head: project against `dec.lm_head.weight` (NOT the - // token embedding). Reference: `MoonshineStreamingForConditionalGeneration.proj_out`. + // Untied lm_head: project against dec.lm_head.weight, NOT the token embed. ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.lm_head_w, x); named(logits_raw, "dec.logits_raw"); if (n_past == 0) { @@ -697,9 +668,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, ggml_set_output(db.out); // Backend argmax over the last position when skipping log_softmax — - // every step pass downloads a single int32 instead of full vocab - // logits. Mirrors cohere/decoder.cpp:756-771 and the moonshine - // sibling change. + // every step pass downloads a single int32 instead of full vocab logits. if (skip_log_softmax) { ggml_tensor * last_logits = logits; if (n_tokens > 1) { @@ -723,10 +692,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } -// =========================================================================== -// Offline batched decode (B utterances) -// =========================================================================== - +// Offline batched decode (B utterances). namespace { ggml_tensor * unpad_head_dim_4d(ggml_context * ctx, ggml_tensor * t, diff --git a/src/arch/moonshine_streaming/encoder.cpp b/src/arch/moonshine_streaming/encoder.cpp index d58d688d..f873e65b 100644 --- a/src/arch/moonshine_streaming/encoder.cpp +++ b/src/arch/moonshine_streaming/encoder.cpp @@ -1,33 +1,15 @@ // arch/moonshine_streaming/encoder.cpp - Moonshine-Streaming encoder // graph builder. // -// Forward shape (streaming-tiny): +// Time-domain frontend: frame reshape [frame_len, T_frames] -> CMVN +// (ggml_norm eps=1e-6, no affine) -> asinh(exp(log_k)·x) -> linear + SiLU -> +// 2× causal Conv1d (left-pad k-1, stride 2; conv1 +bias +SiLU, conv2 +bias +// only). Then n_layers transformer blocks (per-layer sliding-window mask, no +// RoPE) and a final bias-less LN. // -// audio [n_samples=176000] -// -> reshape [frame_len=80, T_frames=2200] (frame_len innermost = ne0) -// -> CMVN per-frame (ggml_norm with eps=1e-6, no affine) -// -> dump enc.embedder.cmvn.out [80, 2200] -// -> asinh(exp(log_k) * x) -// Implemented as: z = x · exp(log_k); y = log(z + sqrt(z² + 1)) -// -> dump enc.embedder.comp.out [80, 2200] -// -> linear (ne=[80, 320] in ggml) → [320, 2200] -// -> SiLU -// -> dump enc.embedder.linear.out [320, 2200] -// -> transpose to [T_frames, hidden] = [2200, 320] for conv_1d -// -> left-pad ne0 by k-1=4 (causal); conv1d stride=2 → [1100, 640] -// -> + bias + SiLU; reshape to [hidden=640, T1=1100] -// -> dump enc.embedder.conv1.out [640, 1100] -// -> transpose to [T1, 640], left-pad 4, conv1d stride=2 → [550, 320] -// -> + bias (NO SiLU); reshape to [hidden=320, T_enc=550] -// -> dump enc.embedder.conv2.out [320, 550] -// -> 6 × transformer block (per-layer sliding-window mask, no RoPE) -// -> final layer_norm (no bias; scale already includes +1.0) -// -> dump enc.final [320, 550] -// -// Encoder LayerNorm note: every `MoonshineStreamingLayerNorm` has -// `unit_offset=True` (effective gain = γ + 1.0). The converter pre-folds -// the +1.0 into the GGUF tensor so we can use the regular `ggml_norm * scale` -// pattern. +// Encoder LayerNorm note: every MoonshineStreamingLayerNorm has +// unit_offset=True (effective gain = γ + 1.0). The converter pre-folds the +// +1.0 into the tensor so we use the plain `ggml_norm * scale` pattern. #include "encoder.h" @@ -285,8 +267,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Audio must be a multiple of frame_len so the reshape is well-defined. // Reference processor right-pads to multiples of 80 with attention_mask - // marking the pad slots. For Stage 4 jfk we use the trimmed length - // and require an exact multiple. + // marking the pad slots; here we require the caller to pass an exact + // multiple (callers right-pad before calling). if (n_samples % frame_len != 0) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "moonshine_streaming encoder: n_samples=%d not a multiple of " diff --git a/src/arch/moonshine_streaming/encoder.h b/src/arch/moonshine_streaming/encoder.h index 272987cf..50339483 100644 --- a/src/arch/moonshine_streaming/encoder.h +++ b/src/arch/moonshine_streaming/encoder.h @@ -4,7 +4,7 @@ // SiLU → 2× causal Conv1d) with N transformer blocks that use // per-layer sliding-window attention masks (no RoPE on the encoder). // -// Dump points (must match `build/validate/moonshine_streaming//dump_coverage.json`): +// Dump points: // enc.audio.in raw PCM input // enc.embedder.cmvn.out CMVN per-frame (eps=1e-6) // enc.embedder.comp.out asinh(exp(log_k) * cmvn_out) @@ -70,11 +70,9 @@ struct EncoderBuild { // Returns 0 if the input is too short. int encoder_t_enc(const MoonshineStreamingHParams & hp, int n_samples); -// Mirror dump_reference_moonshine_streaming_transformers.py::auto_blocks. -// Used by encoder + decoder graph builders to emit debug dumps for the -// same per-layer subset the reference dumper does, so validate.py compares -// matching tensors on both sides instead of generating MISSING-right -// reports for blocks the reference skips. +// Which per-layer blocks the encoder/decoder builders emit as debug dumps. +// Mirrors dump_reference_moonshine_streaming_transformers.py::auto_blocks +// so both sides dump the same subset. bool dump_block_index(int i, int n_layers); // Build the encoder forward graph in compute_ctx. diff --git a/src/arch/moonshine_streaming/model.cpp b/src/arch/moonshine_streaming/model.cpp index c91531cf..2eaad940 100644 --- a/src/arch/moonshine_streaming/model.cpp +++ b/src/arch/moonshine_streaming/model.cpp @@ -1,34 +1,24 @@ // arch/moonshine_streaming/model.cpp - Moonshine-Streaming family handler. // -// Lifecycle: +// Lifecycle: load -> init_context -> run (encoder -> adapter -> cross_kv +// precompute -> autoregressive decode). // -// 1. load — read GGUF, populate hparams, wire weight slots. -// 2. init_context — allocate scheduler / context. -// 3. run — encoder → adapter → cross_kv precompute → -// autoregressive decode. +// One-shot path: the adapter (pos_emb add + optional proj) runs once per +// session in a separate compute_ctx, its output is read back to host, and +// the cross_kv precompute graph projects K/V into the persistent kv_cache. // -// One-shot path: the adapter (pos_emb add + optional proj) runs once -// per session in a separate compute_ctx, its output is read back to -// host, and the cross_kv precompute graph uploads that buffer to -// project K/V projections directly into the persistent kv_cache. -// -// Streaming path: encoder, adapter, and cross-KV K/V projection all -// run incrementally per stream_feed. The encoder is ergodic (no -// positional encoding on encoder self-attn), the adapter's pos_emb is -// a get_rows indexed by absolute frame, and the cross-KV K/V -// projections are per-frame linear, so each stage can run on a slice -// and the slice outputs concatenate across feeds to produce the same -// values a one-shot pass would. Per-feed work is bounded by the slice -// size; the only thing that runs at finalize is one bulk upload of -// the accumulated host K/V buffers into the kv_cache plus the AR -// decoder loop. PCM trimming bounds memory: anything older than -// (T_emitted - L_total - frontend_pad) encoder frames is no longer -// reachable by any future encoder window and is dropped from the -// per-utterance buffer. -// -// AR decoding still happens once at finalize because the model is -// not trained for partial cross-attention; partial transcripts -// mid-stream (re-decode per feed) is a Phase 4b-full follow-up. +// Streaming path: encoder, adapter, and cross-KV projection run incrementally +// per stream_feed, appending each stable slice to host K/V buffers. The +// encoder is ergodic (no positional encoding on encoder self-attn), the +// adapter pos_emb is a get_rows by absolute frame, and the cross-KV +// projection is per-frame linear, so per-feed slice outputs concatenate to +// the one-shot result. A feed can also run a throttled partial AR decode: +// each such decode uploads the accumulated host K/V into the kv_cache and +// decodes from BOS for a live transcript. Finalize tops up the trailing tail +// and runs a final AR decode only if frames advanced past the last partial +// (otherwise it commits the last partial as-is). PCM older than +// (T_emitted - L_total - frontend_pad) encoder frames is unreachable and +// dropped to bound memory. #include "moonshine_streaming.h" @@ -102,25 +92,15 @@ MoonshineStreamingModel::~MoonshineStreamingModel() { plan.primary_kind = transcribe::BackendKind::Unknown; } -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Like moonshine, the wall -// here is on *output*, not input: the encoder accepts any PCM length, and the -// greedy decode loop stops when it reaches the decoder's position cap -// (dec_max_position_embeddings — larger than base moonshine, e.g. 4096) before -// emitting EOS, silently truncating the transcript. We do not gate on audio -// length — a dense clip can exhaust the output budget — so there is no upfront -// INPUT_TOO_LONG rejection. Instead, when a decode loop exits on the cap -// rather than on EOS, we set transcribe_was_truncated() and emit a WARN, same -// convention as qwen3_asr's generation-budget guard. -// --------------------------------------------------------------------------- - -// Advisory transcribe_capabilities::max_audio_ms. The limit is output-bound -// and content-dependent (decode tokens, not audio frames), so there is no -// exact audio ceiling. We report a conservative advisory: the audio duration -// the output budget (dec_max_position_embeddings tokens) can cover at a -// typical ~4 output tokens/sec speaking rate. Returns 0 ("unknown / unbounded") -// if the cap is missing. Crossing this advisory does not reject the clip; it -// only makes a mid-decode truncation (transcribe_was_truncated) more likely. +// 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 int64_t moonshine_streaming_max_audio_ms(const MoonshineStreamingHParams & hp) { if (hp.dec_max_position_embeddings <= 0) { @@ -267,9 +247,8 @@ transcribe_status load( if (auto st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) return st; if (auto st = read_moonshine_streaming_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) return st; - // Publish the advisory input-length window now that the decoder position - // cap is known. apply_family_invariants() ran above (before hparams), so - // this is the first point the cap is available. See docs/input-limits.md. + // 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); gguf_init_params init_params {}; @@ -460,18 +439,9 @@ int samples_per_encoder_frame(const MoonshineStreamingHParams & hp) { return 4 * hp.enc_frame_len; } -// Greedy-decode generation budget, in tokens. Some inputs never trigger EOS -// — a repeated-digit "double nine ... double nine" phone-number clip makes -// the model emit one token forever, identical to the HF reference — so the -// greedy loop needs a bound tighter than the architectural position cap -// (dec_max_position_embeddings, thousands of tokens) to avoid a long, -// pointless compute loop. The upstream model card recommends -// max_new_tokens ~= audio_seconds * 6.5; we follow that, plus a small floor -// so very short clips keep ample headroom, and let the caller take min() -// with the position cap (so this only ever tightens, never loosens, the -// existing wall). Audio length is derived from T_enc: one encoder frame -// spans samples_per_encoder_frame(hp) PCM samples at the 16 kHz native rate. -// Returns 0 when the frame geometry is unknown, meaning "no duration bound". +// Bound greedy decode by duration as well as the decoder position cap. +// Some inputs never emit EOS; upstream recommends about 6.5 tokens/sec. +// Returns 0 when frame geometry is unknown. int decode_generation_budget(const MoonshineStreamingHParams & hp, int T_enc) { if (hp.enc_frame_len <= 0 || T_enc <= 0) { return 0; @@ -487,13 +457,10 @@ int decode_generation_budget(const MoonshineStreamingHParams & hp, int T_enc) { + k_budget_floor); } -// Encoder helper: build the encoder graph for `n_samples` PCM, -// upload PCM + per-layer sliding-window masks, compute, and read the -// final-LN output into the caller-provided host vector. Updates -// cc->t_encode_us (additive). When emit_dumps is true the standard -// encoder.* dump points fire; the streaming feed path passes false so -// per-feed runs don't clobber the reference-dump artifacts that -// validate.py compares. +// 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 +// passes false so per-feed runs do not clobber reference dump artifacts. // // n_samples must be > 0 and a multiple of hp.enc_frame_len. transcribe_status encode_window_to_host( @@ -1066,14 +1033,8 @@ transcribe_status decode_from_kv_cache( } } - // The greedy loop exits either because next_token == eos (complete) or - // because n_past reached gen_cap — the tighter of the duration-based - // generation budget and the hard position cap (max_pos). When the last - // token was not eos the transcript hit the cap before end-of-stream and - // is incomplete: flag it via transcribe_was_truncated() and emit one WARN, - // mirroring qwen3_asr. Reaching the duration budget (gen_cap < max_pos) - // typically means a hallucination loop the model never ends — bounding it - // here avoids a long, pointless decode. See docs/input-limits.md. + // Non-EOS after the loop means gen_cap stopped decode before EOS. gen_cap + // is either the position cap or the tighter duration budget. if (next_token != eos) { cc->was_truncated = true; const bool hit_duration_budget = (gen_cap < max_pos) || (max_pos <= 0); @@ -1279,56 +1240,30 @@ transcribe_status run( if (st != TRANSCRIBE_OK) { return st; } - // Output truncation is a hard status on the OFFLINE single-utterance - // path only (see docs/input-limits.md). decode_from_kv_cache sets - // was_truncated but returns TRANSCRIBE_OK because it is shared with the - // streaming finalize path, which has its own terminal-state machine and - // must NOT surface OUTPUT_TRUNCATED. Scope the remap to this - // transcribe_run entry; the partial transcript stays readable. + // Remap truncation to a hard status only at this offline entry: + // decode_from_kv_cache returns OK (it's shared with the streaming finalize + // path, which must NOT surface OUTPUT_TRUNCATED). Partial text stays readable. return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Streaming hooks -// --------------------------------------------------------------------------- -// -// Per-feed: encode a sliding window, apply adapter on the new emit -// slice, project K/V for every decoder layer, append everything to -// the per-utterance host buffers, then trim PCM that no future -// encoder window will read. Then upload the accumulated host K/V to -// the kv_cache and re-decode from BOS over the current cross-KV so -// the caller sees a live partial transcript. +// Streaming hooks. // -// At finalize: top up the tail (no R_total margin needed since audio -// is done), optionally re-decode if frames advanced since the last -// feed's decode, then mark everything committed. Finalize cost -// bounded by transcript length (just an AR decoder run). +// Per-feed: encode a sliding window, apply adapter + project K/V on the new +// slice, append to host buffers, trim unreachable PCM, then re-decode from BOS +// over the growing cross-KV for a live partial transcript. Finalize tops up +// the tail and runs one last AR decode. // -// PCM trimming retains samples back to -// (T_emitted - L_total - frontend_pad) encoder frames. That window -// covers every future encoder-graph slice (the longest reach is at -// finalize, where the encoder slice may extend back by L_total + -// frontend_pad to seed the conv frontend and self-attn masks). For -// the tiny variant L_total = 90 enc frames = 1.8 s of audio. +// PCM trimming retains samples back to (T_emitted - L_total - frontend_pad) +// encoder frames — the longest reach of any future encoder slice (finalize). +// Tiny: L_total = 90 enc frames = 1.8 s. // -// Commit semantics: -// -// The per-feed decode runs from BOS each time over the growing -// cross-KV. Tokens emitted on feed N may shift on feed N+1 as more -// audio arrives (the model isn't trained for partial cross-attn). -// We mark as committed the longest token-id prefix that's IDENTICAL -// across the last stable_prefix_agreement_n partial decodes (default -// 3). That's the substring the model has now produced consistently -// across multiple cross-KV sizes and is unlikely to revise. Tokens -// past the divergence point are tentative; consumers should treat -// them as preview-only. -// -// n_committed_tokens advances monotonically. n_committed_words / -// n_committed_segments stay at 0 mid-stream (we don't track -// per-word boundaries, and there's exactly one segment per decode). -// At finalize, the full transcript is committed: n_committed_* are -// set to the full counts. +// Commit semantics: per-feed decodes restart from BOS, so feed N's tokens may +// shift on feed N+1 (the model isn't trained for partial cross-attn). The +// committed prefix is the longest token-id prefix IDENTICAL across the last +// stable_prefix_agreement_n decodes (default 3); past that is preview-only. +// n_committed_tokens advances monotonically; n_committed_words/segments stay 0 +// mid-stream and are set to full counts at finalize. constexpr int64_t k_sample_rate_hz = 16000; @@ -1465,22 +1400,13 @@ int stable_token_prefix_from_history( return prefix; } -// Default decode-interval when the caller passes -1 (or omits the -// moonshine_streaming stream params block entirely). Chosen to match -// one cumulative right-context window (R_total = 12 frames × 20 ms), -// giving ~4 partial-transcript updates per second after the initial -// warmup. Compute cost at this setting is roughly 2–3× one-shot on -// the tiny variant; smaller values (e.g. 80 ms) push it to 5×+. +// Default decode-interval (caller passes -1 / omits the stream params block). +// One cumulative right-context window (R_total = 12 frames × 20 ms) = ~4 +// updates/sec; ~2-3× one-shot cost on tiny, smaller values push it to 5×+. constexpr int k_default_min_decode_interval_ms = 240; -// Pure resolver for the Moonshine-Streaming decode-throttle knob. -// Validates the stream_params->family extension's universal shape and -// the min_decode_interval_ms sentinel rule, then resolves the effective -// milliseconds value (family default when -1 / absent). Reads only the -// caller-owned stream_params; mutates nothing. Shared by stream_validate -// (pre-flight, before the dispatcher clears the snapshot) and stream_begin -// (defense in depth, where the resolved value configures session state). -// +// Resolve the decode-throttle knob from the caller-owned stream_params (family +// default when -1 / absent); validates the extension shape. Mutates nothing. // ext NULL / -1 -> family default (k_default_min_decode_interval_ms) // < -1 -> TRANSCRIBE_ERR_INVALID_ARG // >= 0 -> the requested value (0 = decode every advance) @@ -2004,13 +1930,9 @@ bool accepts_ext_kind(const transcribe_model * model, return kind == TRANSCRIBE_EXT_KIND_MOONSHINE_STREAMING_STREAM; } -// =========================================================================== -// Offline batched decode (transcribe_run_batch). Encoder + adapter serial -// per utterance (compute-bound; no mel frontend to parallelize); decode -// (self+cross attn, partial RoPE) batched. Always uses flash (the moonshine -// flash-for-batch decision). See docs/batching-autoregressive-plan.md. -// =========================================================================== - +// Offline batched decode (transcribe_run_batch). Encoder + adapter serial per +// utterance (no mel to parallelize); decode (self+cross attn, partial RoPE) +// batched. Always uses flash (the moonshine flash-for-batch decision). transcribe_status run_batch_serial(MoonshineStreamingSession * cc, const float * const * pcm, const int * n_samples, int n, @@ -2266,12 +2188,9 @@ transcribe_status run_batch( } const int64_t dec_us = ggml_time_us() - t_dec0; - // Batched truncation flag. A valid row that never reached eos exhausted - // the output budget (n_ctx_cap, the position cap clamped to the cache - // capacity) before end-of-stream, so its transcript is incomplete. Mirror - // the serial path: set transcribe_was_truncated() and emit one WARN. This - // only observes the per-row finished[] state; the decode is unchanged. - // See docs/input-limits.md. + // Batched truncation: a valid row that never reached eos exhausted the + // output budget (n_ctx_cap = position cap clamped to cache capacity). + // Mirror the serial path (WARN + flag). { int n_truncated = 0; for (int b = 0; b < n; ++b) { @@ -2304,10 +2223,7 @@ transcribe_status run_batch( rs.full_text = full; rs.result_kind = TRANSCRIBE_TIMESTAMPS_NONE; rs.has_result = true; - // Per-utterance truncation parity (offline run_batch, not streaming): a - // valid row that never finished hit the position cap before EOS, so it - // reports OUTPUT_TRUNCATED (partial text retained). See - // docs/input-limits.md. + // Per-utterance truncation parity (offline run_batch, not streaming). rs.status = !finished[b] ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; rs.t_mel_us = 0; diff --git a/src/arch/moonshine_streaming/moonshine_streaming.h b/src/arch/moonshine_streaming/moonshine_streaming.h index 9c0c39a4..db608a2a 100644 --- a/src/arch/moonshine_streaming/moonshine_streaming.h +++ b/src/arch/moonshine_streaming/moonshine_streaming.h @@ -5,13 +5,8 @@ // src/arch/moonshine/moonshine.h since the encoder-decoder + cross-attn // KV cache shape is shared across both. Differences: // -// - `enc_host_for_adapter` holds the post-adapter encoder hidden -// (instead of the raw encoder output) so the adapter add+proj is -// applied exactly once per session, mirroring the reference dumper's -// `apply_adapter()` helper. Cross-KV precompute reads from this -// buffer, not the encoder output directly. -// - Phase 4a exposes the streaming ABI as stream-of-whole: -// feed buffers PCM and finalize runs the one-shot inference path. +// - The adapter add+proj is applied exactly once per session into a host +// buffer; cross-KV precompute reads that buffer, not the encoder output. #pragma once @@ -140,59 +135,31 @@ struct MoonshineStreamingSession final : public transcribe_session { bool encoder_use_flash = true; bool decoder_use_flash = true; - // ---- incremental streaming state (Phase 4b-encoder + adapter/xkv) ---- + // ---- incremental streaming state ---- // - // Each feed extends three host-side committed buffers in lockstep: + // 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, + // when the throttle allows) and again at + // the finalize decode. // - // stream_adapter_committed - post-adapter encoder hidden, sized - // [dec_d_model, T_emitted]. Built per - // feed by re-running the encoder over - // a sliding window and immediately - // applying the adapter (pos_emb + - // optional proj) at absolute frame - // positions. + // Per-feed slicing is numerically equivalent to a one-shot pass because + // the encoder is ergodic (no positional encoding) with per-layer + // sliding-window attention + causal-conv frontend (output frame t depends + // only on conv-stack frames [t - L_total, t + R_total]), the adapter + // pos_emb is an absolute-frame get_rows, and the cross-KV projection is + // per-frame linear. // - // stream_cross_k_committed - one host vector per decoder layer, - // stream_cross_v_committed each sized [dec_d_model, T_emitted]. - // Populated by the cross-KV projection - // graph (k_proj / v_proj) on the - // adapter slice. NOT written to the - // persistent kv_cache yet — that's a - // single bulk upload at finalize. + // PCM trimming: stream_pcm_buffer drops samples older than + // (T_emitted - L_total - frontend_pad) encoder frames (the left context + // any future window still needs). stream_pcm_start_sample tracks the + // absolute index of stream_pcm_buffer[0]. // - // Trimming: stream_pcm_buffer is trimmed each feed to drop PCM - // older than (T_emitted - L_total - frontend_pad) encoder frames. - // The 1.8 s of left-context history that the encoder needs to - // produce numerically-identical frames is preserved; everything - // before is no longer reachable. stream_pcm_start_sample tracks - // the absolute sample index of stream_pcm_buffer[0] so the - // streaming driver can translate absolute frame indices back to - // buffer-relative offsets. - // - // The design relies on the encoder being ergodic (no positional - // encoding on encoder self-attn) plus per-layer sliding-window - // attention plus causal-conv frontend, so output frame t is a - // function only of conv-stack output frames in - // [t - L_total, t + R_total]; the adapter is a positional - // embedding add indexed by absolute frame so it's free to apply - // per-feed; and the cross-KV projection is per-frame linear so - // accumulating its output across feeds and uploading once at - // finalize is numerically equivalent to a single batched call. - // - // Lifecycle: - // stream_begin: derive geometry, clear scratch + counters. - // stream_feed: append PCM, slide-window encode, apply adapter, - // project cross K/V, append to host buffers, - // trim PCM. - // stream_finalize: encode tail (no R_total margin needed), - // apply adapter + project K/V on the tail, - // allocate kv_cache, upload host K/V buffers, - // run AR decoder only. - // stream_reset: clear scratch (keep capacity). - // - // These fields are NOT touched by clear_result; the dispatcher - // wipes lifecycle-agnostic snapshot state there, and the family - // owns its per-utterance audio + encoder scratch. + // These fields are NOT touched by clear_result — the family owns its + // per-utterance audio + encoder scratch. std::vector stream_pcm_buffer; int64_t stream_pcm_start_sample = 0; std::vector stream_adapter_committed; @@ -212,11 +179,8 @@ struct MoonshineStreamingSession final : public transcribe_session { int32_t stream_R_total_frames = 0; int32_t stream_frontend_pad_frames = 0; int32_t stream_samples_per_enc_frame = 0; - // Minimum gap, in encoder frames, between successive per-feed AR - // decoder runs. Resolved at stream_begin from the Moonshine- - // Streaming family extension. - // 0 means "decode on every encoder-frame advance". The finalize - // path always runs one last decode regardless of this throttle. + // Minimum gap (encoder frames) between per-feed AR decode runs, from the + // family extension. 0 = decode every advance; finalize always decodes once. int32_t stream_min_decode_frames = 0; transcribe_run_params stream_run_params {}; diff --git a/src/arch/parakeet/capabilities.cpp b/src/arch/parakeet/capabilities.cpp index 83a3ef44..e7d20b9a 100644 --- a/src/arch/parakeet/capabilities.cpp +++ b/src/arch/parakeet/capabilities.cpp @@ -1,41 +1,13 @@ // arch/parakeet/capabilities.cpp - Parakeet capability defaults. // -// Per the 2B revision after the variant-dispatch review (see -// RESUME.md "Decisions still load-bearing" -> "stt.variant is -// descriptive metadata, not a behavioral discriminator"): -// -// - The previous variant-driven resolver is gone. Capability -// resolution is now KV-driven via transcribe::read_capability_kv. -// - This file holds only the family-default population, applied -// before the KV reader runs. KV present overrides the family -// default; KV absent leaves the family default in place. PLAN.md -// "GGUF KV is authoritative; if a key is absent, the architecture -// supplies a default" — this function supplies those defaults. -// -// What stays here vs what goes to KV: -// -// - native_sample_rate is set here, not exposed via KV. Parakeet's -// feature extractor is a fixed 16 kHz mel bank for every -// published variant; treating it as a KV-overridable field would -// just create a footgun. If a future variant truly used a -// different rate, the right answer would be to add a separate -// parakeet handler subtree, not to flip a KV. -// -// - supports_translate defaults to false here. Strictly speaking -// the KV reader could overwrite it, but no published Parakeet -// variant has a translate task head and the converter is not -// expected to emit stt.capability.translate for this family. The -// default sticks because nothing overrides it. -// -// - supports_language_detect, supports_streaming come from KV. The -// family default is false; v3 multilingual GGUFs are expected to -// emit stt.capability.lang_detect = true and the KV reader will -// pick it up. -// -// - n_languages / languages come from general.languages via -// read_languages_kv (called by the family handler after this -// function runs). Absent leaves the model in the documented -// "information gap, not a claim" state. +// Family-default capability population, applied before the KV reader +// (transcribe::read_capability_kv) runs: KV present overrides the +// default, KV absent leaves it in place. Variant identity is descriptive +// metadata, not a behavioral discriminator. native_sample_rate is set +// here rather than exposed via KV (fixed 16 kHz mel bank on every +// variant); supports_language_detect / supports_streaming come from KV; +// n_languages / languages come from general.languages via +// read_languages_kv (run by the family handler after this function). #include "parakeet.h" @@ -44,30 +16,19 @@ namespace transcribe::parakeet { void apply_family_invariants(transcribe_model & model) { transcribe_capabilities & caps = model.caps; - // Parakeet's feature extractor is a fixed 16 kHz mel bank (NeMo - // AudioToMelSpectrogramPreprocessor). Every published variant — v2 - // English and v3 multilingual — uses it. The CLI rejects non- - // 16 kHz WAVs at load time so this number is the contract. + // Fixed 16 kHz mel bank (NeMo AudioToMelSpectrogramPreprocessor) on + // every published variant; the CLI rejects non-16 kHz WAVs at load. caps.native_sample_rate = 16000; - // Family default. No published Parakeet variant has a translate - // task head; NeMo ships translation via Canary, not Parakeet. - // KV is authoritative per PLAN.md, so a converter that explicitly - // wrote stt.capability.translate could in principle override - // this — but for Parakeet we never expect that to happen. + // No published Parakeet variant has a translate task head. caps.supports_translate = false; - // Parakeet's TDT decoder emits one token per acoustic frame the - // encoder kept, with step_at_emit carrying the frame index. - // run() converts that to per-token t0/t1 and aggregates up to - // word- and segment-level timings. TOKEN is the finest the - // family can honestly produce, and that is what we advertise. + // TDT decode emits one token per kept encoder frame; run() converts + // step_at_emit to per-token t0/t1 and aggregates to word/segment. + // TOKEN is the finest the family can honestly produce. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_TOKEN; - // Abort callback is honored at the top of each run. Other - // long-form / temperature-fallback / initial-prompt features - // belong to autoregressive families and don't apply here. No - // runtime PNC/ITN control. + // Abort callback honored at the top of each run. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); } diff --git a/src/arch/parakeet/decoder.cpp b/src/arch/parakeet/decoder.cpp index 63aa5342..99524065 100644 --- a/src/arch/parakeet/decoder.cpp +++ b/src/arch/parakeet/decoder.cpp @@ -1,17 +1,15 @@ // arch/parakeet/decoder.cpp - Parakeet TDT decoder implementation. // -// See decoder.h for the API contract. The predictor LSTM, the encoder -// projection, and the joint network all run as ggml backend graphs on a -// single shared CPU backend + persistent threadpool (owned by PredGraph; -// see build_pred_graph). Decode weights are uploaded to resident ggml -// tensors at load time and the host-side fp32 mirrors are then freed. +// See decoder.h for the API contract. The predictor LSTM, encoder +// projection, and joint network all run as ggml backend graphs on a +// single shared CPU backend + persistent threadpool (owned by PredGraph). +// Decode weights are uploaded to resident ggml tensors at load and the +// host fp32 mirrors then freed. // -// Numerical-accuracy strategy: ggml's reduction order differs slightly -// from a naive host loop, so per-step values drift ~1e-7 vs the reference; -// over a full utterance this stays within ~1e-4 and the greedy transcript -// is unchanged. The per-step dump points wired below feed -// scripts/compare_tensors.py for the bring-up loop and are gated on -// TRANSCRIBE_DUMP_DIR. +// Numerical note: ggml's reduction order differs slightly from a naive +// host loop, so per-step values drift ~1e-7 vs the reference; over an +// utterance this stays within ~1e-4 and the greedy transcript is +// unchanged. #include "decoder.h" @@ -22,16 +20,14 @@ #include "transcribe-debug.h" #include "transcribe-log.h" -// ggml-backend.h only — NOT ggml-cpu.h: backend-specific entry points such -// as ggml_backend_cpu_init live inside the loadable CPU module under -// GGML_BACKEND_DL, so the library must reach the CPU backend through the -// registry (ggml_backend_init_by_type / get_proc_address), never by direct -// link-time reference. +// ggml-backend.h, not ggml-cpu.h: under GGML_BACKEND_DL the CPU backend +// is a loadable module, so backend-specific entry points must be reached +// through the registry (ggml_backend_init_by_type / get_proc_address), +// never by direct link-time reference. #include "ggml.h" #include "ggml-backend.h" #include "ggml-alloc.h" -#include "ggml-cpu.h" // ggml_threadpool_params_default (the rest of the CPU - // backend API is reached via the registry, see above) +#include "ggml-cpu.h" // ggml_threadpool_params_default (rest via registry) #include @@ -60,25 +56,11 @@ namespace transcribe::parakeet { namespace { // Read all elements of a ggml_tensor into a host fp32 vector, -// dequantizing as needed via ggml_get_type_traits()->to_float. Decode -// weights are uploaded to resident fp32 ggml tensors once at load time -// (see build_pred_weights / build_joint_weight), so quantized weight -// tensors are dequantized exactly once here and the per-decode hot path -// pays no readback or dequant cost. -// -// Source dtype support is whatever ggml's type traits register a -// to_float implementation for: F32, F16, BF16, all the Q* and IQ* -// quants. Tensors with no to_float fall through to the diagnostic -// (which would be a configuration mistake — the loader's per-slot -// allowlist in weights.cpp should never accept such a type). -// -// ggml_backend_tensor_get is the universal backend API: it's a memcpy -// on host buffers, a readback on discrete GPUs. We stage the raw -// (possibly quantized) bytes into a local buffer first, then walk the -// type's to_float to materialize fp32 in `out`. -// -// Returns false on missing-trait / size errors (logs a diagnostic -// naming the offending tensor). +// dequantizing via ggml_get_type_traits()->to_float. Source dtype +// support is whatever ggml registers a to_float for (F32, F16, BF16, all +// Q*/IQ*). ggml_backend_tensor_get is universal (memcpy on host buffers, +// readback on discrete GPUs); we stage the raw bytes then materialize +// fp32 in `out`. Returns false on missing-trait / size errors. bool read_tensor_to_f32(const ggml_tensor * t, std::vector & out) { @@ -102,10 +84,8 @@ bool read_tensor_to_f32(const ggml_tensor * t, out.resize(static_cast(nelem)); - // F32 fast path. ggml's type_traits[GGML_TYPE_F32] entry does - // NOT set a to_float (the identity case is left null upstream), - // so we must short-circuit here rather than dispatch through it. - // Read straight from the backend into the output buffer. + // F32 fast path. type_traits[GGML_TYPE_F32] has no to_float (identity + // is left null upstream), so short-circuit rather than dispatch. if (t->type == GGML_TYPE_F32) { if (nbytes != static_cast(nelem) * sizeof(float)) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -119,10 +99,8 @@ bool read_tensor_to_f32(const ggml_tensor * t, return true; } - // Quantized / non-fp32 path. Stage the raw bytes off the backend, - // then walk the type's to_float to materialize fp32. F16, BF16, - // and every Q*/IQ* register a to_float in ggml.c's type_traits - // table. + // Quantized / non-fp32 path: stage the raw bytes off the backend, + // then walk the type's to_float to materialize fp32. const auto * tt = ggml_get_type_traits(t->type); if (tt == nullptr || tt->to_float == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -137,15 +115,11 @@ bool read_tensor_to_f32(const ggml_tensor * t, return true; } -// Make the joint network's weights resident as fp32 ggml tensors on the model -// (the immutable half): the encoder projection (g_enc_w/g_enc_b), the predictor -// projection (g_pred_w/g_pred_b), and the output projection (gw_w/gw_b). Built -// once at load; only ever read after, so it is safe to share across every -// context. The per-decode compute graph that consumes these is built per call -// (build_joint_graph below). out_w is dequantized to fp32 from the model tensor -// `src_out_w`; the other weights come from the host mirrors, which are freed -// here once uploaded. On any failure it frees its partial state and returns -// false (w_ready stays false), so the decode reports a hard error. +// Make the joint network's weights resident as fp32 ggml tensors on the +// model: enc / pred / out projections. Built once at load. out_w is +// dequantized from the model tensor src_out_w; the rest come from the +// host mirrors, freed here once uploaded. On failure frees partial state +// and returns false (w_ready stays false → hard decode error). bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w) { const int joint_h = j.joint_h; const int joint_n = j.joint_n; @@ -160,9 +134,8 @@ bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w) { return false; }; - // A CPU backend used ONLY to allocate the resident weight buffer — never - // graph_compute'd (each decode call owns its own compute backend), so it - // carries no per-step state and is safe to keep on the shared model. + // CPU backend used ONLY to allocate the resident weight buffer — + // never graph_compute'd, so it carries no per-step state. j.w_backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); if (j.w_backend == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: ggml CPU backend init failed"); @@ -186,8 +159,7 @@ bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w) { j.w_buf = ggml_backend_alloc_ctx_tensors(j.w_ctx, j.w_backend); if (j.w_buf == nullptr) return fail(); - // out_w: dequantize the model tensor to fp32 once (faithful to the host - // reference: fp32×fp32 accumulation, no activation down-conversion). + // out_w: dequantize the model tensor to fp32 once. { std::vector tmp; if (!read_tensor_to_f32(src_out_w, tmp)) return fail(); @@ -211,12 +183,11 @@ bool build_joint_weight(HostJoint & j, const ggml_tensor * src_out_w) { return true; } -// Make the predictor LSTM weights resident as fp32 ggml tensors — the immutable -// half consumed by the per-call PredGraph. Built once at load. ne fast-to-slow -// is [pred_hidden, 4*pred_hidden] for Wx/Wh (the row-major [4*H, H] host bytes -// read as a ggml mul_mat operand) and [4*pred_hidden] for the bias. On any -// failure frees its partial state and returns false; the caller leaves -// lstm_ready false; building the pred graph then fails and the decode returns a hard error (no host fallback remains). +// Make the predictor LSTM weights resident as fp32 ggml tensors for the +// per-call PredGraph. Built once at load. ne is [pred_hidden, 4*pred_hidden] +// for Wx/Wh (row-major [4*H, H] host bytes as a mul_mat operand) and +// [4*pred_hidden] for the bias. On failure frees partial state and returns +// false (lstm_ready stays false → hard decode error). bool build_pred_weights(HostPredictor & p) { const int H = p.pred_hidden; const int four_H = 4 * H; @@ -261,7 +232,7 @@ bool build_pred_weights(HostPredictor & p) { ggml_backend_tensor_set(lh.g_Wx, lh.Wx.data(), 0, lh.Wx.size() * sizeof(float)); ggml_backend_tensor_set(lh.g_Wh, lh.Wh.data(), 0, lh.Wh.size() * sizeof(float)); ggml_backend_tensor_set(lh.g_b, lh.b.data(), 0, lh.b.size() * sizeof(float)); - // Host mirrors are now resident in ggml — release them (load-time scratch). + // Host mirrors now resident in ggml — release them. std::vector().swap(lh.Wx); std::vector().swap(lh.Wh); std::vector().swap(lh.b); @@ -273,22 +244,16 @@ bool build_pred_weights(HostPredictor & p) { } // namespace -// --------------------------------------------------------------------------- -// Per-call joint compute graph -// --------------------------------------------------------------------------- -// -// The MUTABLE half of the joint network: a full per-call graph +// Per-call joint compute graph (the mutable half of the joint network): // pred_proj = pred_w @ pred_in + pred_b [joint_h] // summed = enc_proj + pred_proj [joint_h] // activated = activation(summed) [joint_h] // logits = out_w @ activated + out_b [joint_n] -// built on a BORROWED backend — the shared decoder pool owned by PredGraph — so -// the joint runs on the SAME single threadpool as the pred LSTM and enc_proj -// (no separate pool, no oversubscription). Inputs are the predictor output -// [pred_hidden] and this frame's precomputed encoder projection [joint_h]; -// output is the logits [joint_n]. Built fresh per decode call (reentrant: each -// decode owns its own I/O tensors); weights are the shared resident -// HostJoint::g_pred_w/g_pred_b/gw_w/gw_b. +// Built on a BORROWED backend (PredGraph's shared pool) so the joint runs +// on the same threadpool as the pred LSTM and enc_proj. Inputs are the +// predictor output and this frame's precomputed enc projection; output is +// the logits. Built fresh per decode call (reentrant); weights are the +// shared resident HostJoint tensors. namespace { struct JointGraph { @@ -305,17 +270,15 @@ struct JointGraph { ~JointGraph() { if (buf != nullptr) ggml_backend_buffer_free(buf); if (ctx != nullptr) ggml_free(ctx); - // backend is borrowed from PredGraph — do NOT free here (and PredGraph, - // declared first in the driver, is destroyed after this). + // backend is borrowed from PredGraph — do NOT free here. } JointGraph(const JointGraph &) = delete; JointGraph & operator=(const JointGraph &) = delete; }; -// Build the full joint graph on the shared `backend` (PredGraph's pool). Returns -// false (g.ready == false) if the resident weights are absent or any ggml step -// fails; the driver treats a non-ready joint graph as a hard decode error (there -// is no host fallback — the resident weights are the only copy). +// Build the full joint graph on the shared `backend` (PredGraph's pool). +// Returns false if the resident weights are absent or any ggml step +// fails (→ hard decode error). bool build_joint_graph(JointGraph & g, const HostJoint & j, ggml_backend_t backend) { if (backend == nullptr) return false; if (!j.w_ready || j.gw_w == nullptr || j.gw_b == nullptr || @@ -371,11 +334,9 @@ bool build_joint_graph(JointGraph & g, const HostJoint & j, ggml_backend_t backe return true; } -// CPU-backend threadpool entry points, reached through the registry so the -// library stays DL-safe: under GGML_BACKEND_DL the CPU backend is a loadable -// module and these symbols are NOT linkable directly into libtranscribe (only -// ggml-base symbols like ggml_threadpool_params_default are). Resolve them via -// ggml_backend_reg_get_proc_address, exactly as for ggml_backend_set_n_threads. +// CPU-backend threadpool entry points, reached through the registry so +// the library stays DL-safe (under GGML_BACKEND_DL these symbols are not +// directly linkable). Resolved via ggml_backend_reg_get_proc_address. typedef ggml_threadpool_t (*pfn_threadpool_new)(ggml_threadpool_params *); typedef void (*pfn_threadpool_free)(ggml_threadpool_t); typedef void (*pfn_set_threadpool)(ggml_backend_t, ggml_threadpool_t); @@ -395,18 +356,12 @@ static void free_cpu_threadpool(ggml_backend_t backend, ggml_threadpool_t tp) { } } -// --------------------------------------------------------------------------- -// Per-call predictor (LSTM) compute graph -// --------------------------------------------------------------------------- -// -// The MUTABLE half of the predictor: a single per-step LSTM graph built fresh -// per decode call around the model-resident HostPredictor::lstm weights, and -// recomputed in place each step. Inputs are the embedding x and, per layer, the -// previous (h, c); outputs are the new (h, c) per layer. The decode loop feeds -// prev state in and reads new state back into its host LstmState each step, so -// the loop's state machine (swap / predictor_dirty cache / dumps) is unchanged. -// Mirrors JointGraph's reentrancy: every concurrent decode owns its own -// PredGraph (its own backend + threadpool + I/O tensors). +// Per-call predictor LSTM graph (the mutable half of the predictor): +// a single per-step graph built fresh per decode call around the +// model-resident HostPredictor::lstm weights, recomputed in place each +// step. Inputs are the embedding x and per-layer previous (h, c); outputs +// are the new (h, c) per layer. Reentrant: every concurrent decode owns +// its own PredGraph (backend + threadpool + I/O tensors). struct PredGraph { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -434,9 +389,8 @@ struct PredGraph { }; // Build the per-call LSTM graph around the resident p.lstm[*].g_Wx/g_Wh/g_b. -// Returns false (g.ready == false) if the resident weights are absent or any -// ggml step fails; the driver then returns a hard decode error (no host fallback). n_threads is -// the resolved (>0) decode thread count. +// Returns false if the resident weights are absent or any ggml step fails +// (→ hard decode error). n_threads is the resolved (>0) thread count. bool build_pred_graph(PredGraph & g, const HostPredictor & p, int n_threads) { if (!p.lstm_ready) return false; const int H = p.pred_hidden; @@ -553,11 +507,8 @@ HostPredictor::~HostPredictor() { if (lstm_w_backend != nullptr) ggml_backend_free(lstm_w_backend); } -// Resolve a decode thread count: n_threads <= 0 means "auto" → min(8, usable -// cpus) via default_n_threads(), matching the encoder default. Threaded through -// to the joint compute backend -// and the per-call ggml decode graphs; no process-global state is set, -// so concurrent decodes on one model do not stomp each other. +// Resolve a decode thread count: n_threads <= 0 means "auto" → +// default_n_threads() (min(8, usable cpus)), matching the encoder. static int resolve_decode_threads(int n_threads) { return n_threads > 0 ? n_threads @@ -579,13 +530,10 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, } if (hp.head_kind == HeadKind::CTC) { - // ----- CTC head mirror ----- - // - // Source weight has ggml ne [1, d_model, n_classes] (PyTorch - // [n_classes, d_model, 1] flat). The bytes are already laid out - // exactly the way the host decoder wants to consume them - // (row-major [n_classes, d_model]); read_tensor_to_f32 just - // copies them into a flat fp32 vector. Bias is [n_classes]. + // CTC head mirror. Source weight ne [1, d_model, n_classes] + // (PyTorch [n_classes, d_model, 1]); bytes are already row-major + // [n_classes, d_model] as the host decoder consumes them. Bias is + // [n_classes]. out.ctc_head.n_classes = hp.head_ctc_n_classes; out.ctc_head.blank_id = hp.head_ctc_n_classes - 1; // NeMo convention out.ctc_head.d_enc = hp.enc_d_model; @@ -629,10 +577,7 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, if (!read_tensor_to_f32(w.predictor.embed_w, out.predictor.embed_w)) { return TRANSCRIBE_ERR_GGUF; } - // Sanity-check the flattened size against the catalog shape. The - // loader already validated ne against hparams; this is a cheap - // belt-and-braces in case future surgery to either side gets out - // of sync without updating the other. + // Sanity-check the flattened size against the catalog shape. { const size_t expected = static_cast(hp.pred_vocab) * static_cast(hp.pred_hidden); @@ -667,9 +612,8 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, } } - // Make the predictor LSTM weights resident as fp32 ggml tensors for the ggml - // predictor graph. Fatal: there is no host fallback, so a failure here means - // the model cannot decode — fail fast at load rather than at first decode. + // Make the predictor LSTM weights resident. Fatal: a failure means + // the model cannot decode — fail fast at load. if (!build_pred_weights(out.predictor)) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: predictor ggml weight build failed"); @@ -683,20 +627,16 @@ transcribe_status build_host_decoder_weights(const ParakeetModel & model, out.joint.joint_n = hp.joint_n_classes(); out.joint.activation = hp.joint_activation; - // enc/pred/out_b are mirrored to host fp32 (build_joint_weight uploads them - // into resident ggml tensors and frees them). out_w is NOT mirrored here: - // build_joint_weight dequantizes it straight from the model tensor. + // enc/pred/out_b are mirrored to host fp32; out_w is NOT mirrored + // here (build_joint_weight dequantizes it from the model tensor). if (!read_tensor_to_f32(w.joint.enc_w, out.joint.enc_w)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.enc_b, out.joint.enc_b)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.pred_w, out.joint.pred_w)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.pred_b, out.joint.pred_b)) return TRANSCRIBE_ERR_GGUF; if (!read_tensor_to_f32(w.joint.out_b, out.joint.out_b)) return TRANSCRIBE_ERR_GGUF; - // Make the joint weights resident as fp32 ggml tensors for the joint graph. - // fp32 keeps decode numerically faithful to the historical host path (the - // out_w matmul was always fp32 there); any joint quantization is handled at - // the model/quant level. Fatal: with no host fallback, a failure here means - // the model cannot decode — fail fast at load. + // Make the joint weights resident as fp32 ggml tensors. Fatal: a + // failure means the model cannot decode — fail fast at load. if (w.joint.out_w == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder: joint out_w missing"); return TRANSCRIBE_ERR_GGUF; @@ -744,13 +684,11 @@ void LstmState::reset(int n_layers, int pred_hidden) { namespace { -// ggml-graph variant of predictor_step. Same contract: -// reads `prev_state`, writes the new step's (h, c) into `new_state`, returns a -// borrowed pointer into `new_state.h.back()`. Internally it feeds prev state and -// the embedding into the resident per-call graph, computes, and reads the new -// state back into the host LstmState — so the decode loop's state machine (swap, -// predictor_dirty cache, dumps) is identical to the host path. Caller guarantees -// g.ready and that new_state is sized to (L, H). +// One predictor step on the ggml graph: reads prev_state, writes the new +// step's (h, c) into new_state, returns a borrowed pointer into +// new_state.h.back(). Feeds prev state + embedding into the resident +// per-call graph, computes, reads new state back into the host LstmState. +// Caller guarantees g.ready and that new_state is sized to (L, H). const float * predictor_step_ggml(const HostPredictor & predictor, PredGraph & g, int last_token, @@ -763,7 +701,7 @@ const float * predictor_step_ggml(const HostPredictor & predictor, scratch_x.resize(static_cast(H)); } - // Embed lookup (or start-state zeros) — identical to predictor_step. + // Embed lookup (or start-state zeros). if (last_token < 0) { std::fill_n(scratch_x.data(), H, 0.0f); } else { @@ -793,19 +731,14 @@ const float * predictor_step_ggml(const HostPredictor & predictor, // Run the joint network for one decode step. // // `enc_proj` is the precomputed encoder projection for this frame -// (enc_w @ enc_frame + enc_b, `joint_h` floats). The caller batches -// all T_enc projections via one ggml GEMM (precompute_enc_proj_ggml) -// before the decode loop so they are computed once rather than -// redundantly per iteration. -// +// (enc_w @ enc_frame + enc_b, joint_h floats); the caller batches all +// T_enc projections via one GEMM (precompute_enc_proj_ggml) before the +// loop. The graph computes: // pred_proj = pred_w @ pred_state + pred_b [joint_h] // summed = enc_proj + pred_proj [joint_h] // activated = activation(summed) [joint_h] // logits = out_w @ activated + out_b [joint_n] -// -// Activation is one of {relu, sigmoid, tanh}; the loader's -// allow-list guarantees we'll see exactly one of those at run -// time. Parakeet 0.6B v2/v3 ship "relu". +// Activation is one of {relu, sigmoid, tanh} (loader allow-list). void joint_step(const HostJoint & j, const JointGraph & g, const float * enc_proj, @@ -814,10 +747,7 @@ void joint_step(const HostJoint & j, { if (static_cast(out_logits.size()) < j.joint_n) out_logits.resize(static_cast(j.joint_n)); - // Full joint on the shared decoder pool: pred_w proj + sum + activation + - // out_w proj + bias, one graph, one dispatch. Inputs are the predictor - // output and this frame's enc_proj; the weights (incl. activation) are baked - // into the graph. + // Full joint on the shared decoder pool, one graph, one dispatch. ggml_backend_tensor_set(g.pred_in, pred_state, 0, static_cast(j.pred_hidden) * sizeof(float)); ggml_backend_tensor_set(g.enc_in, enc_proj, 0, @@ -826,34 +756,13 @@ void joint_step(const HostJoint & j, ggml_backend_tensor_get(g.logits, out_logits.data(), 0, static_cast(j.joint_n) * sizeof(float)); - // Apply log_softmax over the full joint output (vocab+blank+durations - // for TDT, vocab+blank for RNNT) to match NeMo's CPU-inference - // representation of `joint_after_projection`. NeMo's RNNTJoint applies - // `log_softmax(dim=-1)` when running on CPU and `self.log_softmax` - // is None (the default), or when `self.log_softmax=True`. Applying - // it here: - // - leaves token-argmax and duration-argmax invariant (subtracting - // a constant from all elements of either sub-range preserves - // argmax), - // - leaves token_confidence invariant (it re-softmaxes the token - // sub-range, which absorbs any uniform shift), - // - lets `dec.joint.0` dump compare element-wise against NeMo's - // reference dump (representation match — without this, our raw - // logits sit ~+25 to +35 above NeMo's log-softmax-normalized - // values, and the only flat tolerance that fits is 100/100, - // which is loose enough to hide real numerical divergence). - // - // Cost: ~joint_n adds + 1 logsumexp. For joint_n=1030 that's ~2K - // ops per decode iter, negligible vs the joint_h × pred_hidden + - // joint_h × joint_n matmuls (~1.3M ops). - // - // GREEDY FAST PATH: the log_softmax is a uniform per-row shift, so it leaves - // BOTH the token/duration argmax AND token_confidence (which re-softmaxes the - // token sub-range, absorbing the shift) invariant. It is therefore only - // needed to make the `dec.joint.0` dump comparable to NeMo's normalized - // reference. Skip it entirely unless dumping — saves a joint_n-wide - // max+exp+log+sub pass (joint_n≈1030, with a double exp) on every decode - // iteration. Bit-identical decode output either way. + // log_softmax over the full joint output matches NeMo's CPU-inference + // representation of joint_after_projection. It is a uniform per-row + // shift, so it leaves token/duration argmax AND token_confidence + // (which re-softmaxes the token sub-range) invariant — the decode + // output is bit-identical with or without it. It is only needed to + // make the dec.joint.0 dump element-wise comparable to NeMo's + // normalized reference, so skip it unless dumping. if (transcribe::debug::enabled()) { float max_v = out_logits[0]; for (int i = 1; i < j.joint_n; ++i) { @@ -871,13 +780,11 @@ void joint_step(const HostJoint & j, } // Compute the per-utterance encoder projection out[T, joint_h] = -// enc_out[T, d_enc] @ enc_w^T + enc_b as a single GEMM on `backend` — which -// carries the shared decoder threadpool, so this runs on the same pool as the -// pred/joint graphs (no extra pool). The weight + bias are the model-resident -// j.g_enc_w / j.g_enc_b (no per-call upload); only the activation input and the -// result are allocated here. At n = T the matmul is a real GEMM, so with -// GGML_LLAMAFILE=ON tinyBLAS engages (n >= 2). Returns false on any ggml -// failure; the driver treats it as a hard decode error (no host fallback). +// enc_out[T, d_enc] @ enc_w^T + enc_b as a single GEMM on `backend` +// (the shared decoder pool). Weight + bias are the model-resident +// j.g_enc_w / j.g_enc_b; only the input and result are allocated here. +// At n = T the matmul is a real GEMM (tinyBLAS engages with +// GGML_LLAMAFILE=ON). Returns false on ggml failure (→ hard decode error). bool precompute_enc_proj_ggml(const HostJoint & j, ggml_backend_t backend, const float * enc_out, @@ -1014,15 +921,11 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, const int n_dur = static_cast(w.tdt_durations.size()); const int blank_id = w.blank_id; - // Per-call joint compute graph around the shared resident weight. Owns its - // own backend + I/O tensors, so concurrent decodes don't share mutable - // state. A graph build failure is a hard decode error (guarded at driver entry). - // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + - // pool; enc_proj and the joint graph borrow it. pg is declared first so it - // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode - // call, reentrant. The CPU backend is always available, so these succeed in - // practice; a build failure is a hard error (the host decode path was - // removed when the migration finalized). + // All-ggml decode on ONE shared threadpool: PredGraph owns the + // backend + pool; enc_proj and the joint graph borrow it. pg is + // declared first so it is destroyed LAST (jg's buffer lives on pg's + // backend). Built per decode call, reentrant. A build failure is a + // hard decode error. PredGraph pg; JointGraph jg; build_pred_graph(pg, w.predictor, nt); @@ -1063,25 +966,16 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, int step = 0; int new_symbols = 0; - // Honest cap: a 30-second clip emits a few hundred tokens at - // most. Bound the loop iterations as a runaway-protection net - // in case the joint produces pathological logits we haven't - // anticipated. The cap is generous; legitimate transcriptions - // never approach it. + // Runaway-protection cap; legitimate transcriptions never approach it. const int max_iters = 16 * T_enc + 1024; int iter = 0; int64_t t_pred_us = 0, t_joint_us = 0, t_conf_us = 0; - // On a blank emission `(last_token, state)` are preserved, so the - // predictor would deterministically write the same `next_state` on the - // following iteration. Skip the LSTM unroll and reuse the previous - // `decoder_out`; recompute on a non-blank emission, where - // `state.swap(next_state)` and `last_token` change. Bit-exact with the - // unconditional path. Complements the `duration==0` fast-forward - // below: that fast-forward only fires when `tdt_max_symbols > 0` and - // duration is zero. Blanks with `duration > 0` advance `step` but - // still leave `(last_token, state)` unchanged, and this cache catches - // them. + // On a blank emission (last_token, state) are preserved, so the + // predictor would write the same next_state next iteration. Skip the + // LSTM unroll and reuse the previous decoder_out; recompute on a + // non-blank emission (state.swap + last_token change). Bit-exact with + // the unconditional path. bool predictor_dirty = true; while (step < T_enc && iter < max_iters) { @@ -1116,16 +1010,9 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, const int duration = w.tdt_durations[static_cast(decision)]; - // ----- Optional dump (first iteration only, for bring-up) ----- - // - // Dumping every iteration would flood the dump dir on a long - // clip; the bring-up only needs a few well-chosen sample - // points to verify each component matches the NeMo reference. - // The first decode step (start state, encoder frame 0) gives - // us four reference points: predictor input (the embed lookup, - // here all zeros), per-layer LSTM h, the joint logits, the - // argmax decision. The Python `dump_reference_parakeet_nemo.py - // decode` subcommand mirrors these. + // Optional dump of the first decode step (start state, encoder + // frame 0): embed input, per-layer LSTM h, joint logits. Mirrors + // the Python reference dumper's decode subcommand. if (iter == 1 && transcribe::debug::enabled()) { // Predictor scratch_x at start: zeros (vector of length H). const long long s_h = H; @@ -1266,13 +1153,7 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const int n_token_cls = w.predictor.pred_vocab; // == vocab + 1 const int blank_id = w.blank_id; - // Per-call joint compute graph (see decode_tdt_greedy). - // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + - // pool; enc_proj and the joint graph borrow it. pg is declared first so it - // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode - // call, reentrant. The CPU backend is always available, so these succeed in - // practice; a build failure is a hard error (the host decode path was - // removed when the migration finalized). + // Per-call decode graphs (see decode_tdt_greedy for the lifecycle). PredGraph pg; JointGraph jg; build_pred_graph(pg, w.predictor, nt); @@ -1310,14 +1191,9 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, int iter = 0; int64_t t_pred_us = 0, t_joint_us = 0, t_conf_us = 0; - // On a blank emission `(last_token, state)` are preserved, so the - // predictor would deterministically write the same `next_state` on the - // following iteration. Skip the LSTM unroll and reuse the previous - // `decoder_out`; recompute on a non-blank emission, where - // `state.swap(next_state)` and `last_token` change. Bit-exact with the - // unconditional path. On a 0.6B FastConformer RNN-T with a typical - // 4-5× iters/token ratio this elides 60-80% of predictor work and - // ~halves decode wall time. + // Predictor-step cache (see decode_tdt_greedy). Bit-exact with the + // unconditional path; elides 60-80% of predictor work on a typical + // RNN-T iters/token ratio. bool predictor_dirty = true; while (step < T_enc && iter < max_iters) { @@ -1438,14 +1314,9 @@ transcribe_status decode_rnnt_greedy_streaming( if (enc_out == nullptr || T_enc_new <= 0 || d_enc <= 0) { return TRANSCRIBE_ERR_INVALID_ARG; } - // This driver implements pure RNN-T greedy search: every non-blank - // advances the predictor by one symbol and every blank advances the - // encoder frame by exactly one. A TDT head's duration predictions - // would be silently ignored here (collapsing every blank to a single - // frame) and a CTC head carries no predictor state at all. Both of - // today's streaming parakeet variants are RNN-T, so this only guards - // a hypothetical future TDT/CTC streaming model — but it fails loud at - // the exact point of mis-decode rather than emitting wrong tokens. + // Pure RNN-T greedy only: a TDT head's durations would be silently + // ignored and a CTC head carries no predictor state. Fail loud rather + // than mis-decode a hypothetical future TDT/CTC streaming model. if (w.head_kind != HostHeadKind::RNNT) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet decoder (rnnt-stream): streaming decode " @@ -1466,13 +1337,7 @@ transcribe_status decode_rnnt_greedy_streaming( const int n_token_cls = w.predictor.pred_vocab; const int blank_id = w.blank_id; - // Per-call joint compute graph (see decode_tdt_greedy). - // All-ggml decode on ONE shared threadpool: PredGraph owns the backend + - // pool; enc_proj and the joint graph borrow it. pg is declared first so it - // is destroyed LAST (jg's buffer lives on pg's backend). Built per decode - // call, reentrant. The CPU backend is always available, so these succeed in - // practice; a build failure is a hard error (the host decode path was - // removed when the migration finalized). + // Per-call decode graphs (see decode_tdt_greedy for the lifecycle). PredGraph pg; JointGraph jg; build_pred_graph(pg, w.predictor, nt); @@ -1483,7 +1348,7 @@ transcribe_status decode_rnnt_greedy_streaming( return TRANSCRIBE_ERR_BACKEND; } - // Validate state_io shape; reset if degenerate (paranoid guard). + // Validate state_io shape; reset if degenerate. if (static_cast(state_io.h.size()) != n_layers || static_cast(state_io.c.size()) != n_layers) { @@ -1508,10 +1373,9 @@ transcribe_status decode_rnnt_greedy_streaming( return TRANSCRIBE_ERR_BACKEND; } - // Working state. We mutate `state` and `next_state`; on return we - // copy `state` (the COMMITTED state after the last non-blank - // emission, or unchanged for an all-blank chunk) back into - // state_io. last_token is committed similarly. + // Working state. On return, `state` (committed after the last + // non-blank emission, or unchanged for an all-blank chunk) and + // last_token are copied back into state_io / last_token_io. LstmState state = state_io; LstmState next_state; next_state.reset(n_layers, H); @@ -1588,19 +1452,11 @@ transcribe_status decode_rnnt_greedy_streaming( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// CTC greedy decode -// --------------------------------------------------------------------------- +// CTC greedy decode. // -// Per-frame: logits[t] = W @ enc[t] + b -> log_softmax -> argmax. -// Collapse: drop adjacent duplicates ("aaab" -> "ab"), then drop the -// blank label. Standard CTC greedy (see NeMo GreedyCTCInfer). -// -// Dump points (gated on TRANSCRIBE_DUMP_DIR): -// - dec.ctc.logprobs : full [T_enc, n_classes] log-softmax -// - dec.ctc.logprobs.0 : frame 0 of the same -// These match the Stage 2 oracle's reference dumps. - +// Per-frame: logits[t] = W @ enc[t] + b -> log_softmax -> argmax. +// Collapse: drop adjacent duplicates ("aaab" -> "ab"), then drop blanks +// (standard CTC greedy, cf. NeMo GreedyCTCInfer). transcribe_status decode_ctc_greedy(const HostDecoderWeights & w, const float * enc_out, int T_enc, @@ -1633,10 +1489,9 @@ transcribe_status decode_ctc_greedy(const HostDecoderWeights & w, w.ctc_head.weight.data(), d_enc, 0.0f, logits_all.data(), n_classes); #else - // No BLAS: project all T frames in parallel over the shared parallel_for_all - // helper (one thread spawn for the utterance; rows within a frame are a - // serial, auto-vectorized dot). The bias is folded in below. (nt is unused - // on BLAS builds — cblas_sgemm owns its own threading — so it lives here.) + // No BLAS: project all T frames in parallel over parallel_for_all + // (rows within a frame are a serial, auto-vectorized dot). Bias folded + // in below. { const int nt = resolve_decode_threads(n_threads); const float * Wc = w.ctc_head.weight.data(); diff --git a/src/arch/parakeet/decoder.h b/src/arch/parakeet/decoder.h index 04de1cc9..c3b50557 100644 --- a/src/arch/parakeet/decoder.h +++ b/src/arch/parakeet/decoder.h @@ -1,37 +1,29 @@ // arch/parakeet/decoder.h - Parakeet TDT decoder. // -// Implements the predictor (2-layer LSTM) + joint network forward + -// greedy decode driver (TDT, RNN-T, and CTC heads). +// Predictor (2-layer LSTM) + joint network forward + greedy decode +// driver (TDT, RNN-T, CTC heads). // -// Execution model. The decoder runs entirely as ggml graphs on ONE shared, -// per-decode CPU threadpool: the predictor LSTM (per-step graph), the -// per-utterance encoder projection (one GEMM), and the joint network (pred_w -// proj + sum + activation + out_w proj, one graph). PredGraph owns the backend -// and threadpool; the enc_proj and joint graphs borrow it, so the whole decode -// shares a single pool with no oversubscription. Every decode call builds its -// own graphs (reentrant); the weights are model-resident fp32 ggml tensors -// (HostPredictor::lstm_w_* and HostJoint::gw_*/g_pred_*), built once at load. +// Execution model: the decoder runs entirely as ggml graphs on ONE +// shared, per-decode CPU threadpool — predictor LSTM (per-step graph), +// per-utterance encoder projection (one GEMM), and joint network (one +// graph). PredGraph owns the backend and threadpool; enc_proj and joint +// borrow it, so the whole decode shares a single pool with no +// oversubscription. Every decode call builds its own graphs (reentrant); +// weights are model-resident fp32 ggml tensors built once at load. The +// CPU backend always exists (the decoder runs on host even with the +// model on a GPU), so a graph build failure is a hard decode error. // -// This is a deliberate "one threading system" design: the decoder used to run -// host fp32 matmuls on a bespoke worker pool, which coexisted badly with ggml's -// own pool (teardown / oversubscription / MSVC-vcomp bugs). Routing everything -// through ggml deleted that second system. The CPU backend is always available -// (the decoder runs on host even when the model is on a GPU), so a graph build -// failure is a hard decode error rather than a host fallback. +// All weights are fp32: the LSTM and joint matmuls are n=1 GEMV per step +// where quantization buys no usable bandwidth and only adds drift. Any +// joint quantization is a model/quant-level concern. // -// All weights are fp32 (the LSTM and joint matmuls are n=1 GEMV per step, where -// quantization buys no usable bandwidth and would only add drift); fp32 keeps -// decode numerically faithful to the historical host reference. Any joint -// quantization is a model/quant-level concern, not done here. +// Memory cost: host + resident-ggml mirror of predictor + joint weights +// (~35 MB v2, ~73 MB v3) vs the ~2.4 GB encoder. Built once in +// build_host_decoder_weights via ggml_backend_tensor_get (universal +// across host buffers, Metal unified memory, discrete GPUs). // -// Memory cost: a load-time host + resident-ggml mirror of the predictor + joint -// weights (~35 MB on v2, ~73 MB on v3) against the ~2.4 GB encoder footprint. -// The mirror is built once in build_host_decoder_weights via -// ggml_backend_tensor_get (universal across host buffers, Metal unified -// memory, and discrete GPUs). -// -// This header is INTERNAL to src/arch/parakeet/. The public C ABI -// only sees ParakeetSession::result populated by Parakeet::run. +// Internal to src/arch/parakeet/; the public C ABI only sees +// ParakeetSession::result populated by Parakeet::run. #pragma once @@ -67,16 +59,15 @@ struct ParakeetHParams; // W_ih + W_hh bias pre-summed at conversion time. struct HostLstmLayer { - // Host fp32 mirrors — LOAD-TIME SCRATCH ONLY: filled at load, uploaded into - // the resident ggml tensors below by build_pred_weights, then freed. + // Host fp32 mirrors — LOAD-TIME SCRATCH ONLY: filled at load, + // uploaded into the resident ggml tensors below, then freed. std::vector Wx; // [4*pred_hidden, pred_hidden] (freed after load) std::vector Wh; // [4*pred_hidden, pred_hidden] (freed after load) std::vector b; // [4*pred_hidden] (freed after load) - // Resident fp32 ggml mirrors of the three weights above, consumed by the - // ggml predictor graph. Row-major [4*H, H] host bytes map to ggml ne - // fast-to-slow [H, 4*H] (the mul_mat operand); bias is [4*H]. Borrowed - // pointers into HostPredictor::lstm_w_ctx — owned/freed by ~HostPredictor. + // Resident fp32 ggml mirrors consumed by the predictor graph. + // Row-major [4*H, H] host bytes map to ggml ne [H, 4*H] (the mul_mat + // operand); bias is [4*H]. Borrowed into lstm_w_ctx, freed by dtor. ggml_tensor * g_Wx = nullptr; ggml_tensor * g_Wh = nullptr; ggml_tensor * g_b = nullptr; @@ -89,14 +80,11 @@ struct HostPredictor { std::vector lstm; // pred_n_layers entries // --- resident ggml LSTM weights (immutable, model-owned) --- - // fp32 mirror of the per-layer Wx/Wh/b, made resident once at load so the - // per-decode-call PredGraph reads them without re-uploading. Same - // immutable-half / mutable-per-call split as HostJoint::gw_* + the - // stack-local JointGraph. fp32 (not native quant): the per-step LSTM - // matmuls are n=1 GEMV — tinyBLAS skips them (n<2) and they gain no usable - // weight-bandwidth benefit, so fp32 stays closest to the host reference at - // no speed cost. Owned here; freed by ~HostPredictor. On build failure - // lstm_ready stays false and the decode reports a hard error. + // fp32 mirror of the per-layer Wx/Wh/b, made resident once at load so + // the per-decode PredGraph reads them without re-uploading. fp32: the + // per-step LSTM matmuls are n=1 GEMV (tinyBLAS skips n<2), so fp32 is + // closest to the host reference at no speed cost. Owned here, freed by + // dtor. On build failure lstm_ready stays false (hard decode error). ggml_context * lstm_w_ctx = nullptr; ggml_backend_t lstm_w_backend = nullptr; // alloc-only; never compute'd ggml_backend_buffer_t lstm_w_buf = nullptr; @@ -117,11 +105,9 @@ struct HostJoint { int joint_n = 0; // total output classes (vocab+blank+durations) std::string activation; // "relu" / "sigmoid" / "tanh" - // Host fp32 weight mirrors — LOAD-TIME SCRATCH ONLY. read_tensor_to_f32 - // fills these at load; build_joint_weight uploads them into the resident - // ggml tensors below and then frees them, so the decode hot path reads only - // the resident tensors. (out_w is not mirrored here — gw_w is built straight - // from the model tensor.) + // Host fp32 weight mirrors — LOAD-TIME SCRATCH ONLY: filled at load, + // uploaded into the resident ggml tensors below, then freed. (out_w is + // not mirrored — gw_w is built straight from the model tensor.) std::vector enc_w; // [joint_h, d_enc] (freed after load) std::vector enc_b; // [joint_h] (freed after load) std::vector pred_w; // [joint_h, pred_hidden] (freed after load) @@ -129,15 +115,12 @@ struct HostJoint { std::vector out_b; // [joint_n] (freed after load) // --- resident ggml weights (immutable, model-owned) --- - // The whole joint runs as one ggml graph (build_joint_graph), so every joint - // weight is resident as an fp32 ggml tensor: the encoder projection - // (g_enc_w/g_enc_b), the predictor projection (g_pred_w/g_pred_b), and the - // output projection (gw_w/gw_b). fp32 keeps decode faithful to the host - // reference; any joint quantization is a model/quant-level concern. Built - // once at load and only read after — safe to share across every context; - // the mutable per-decode state lives in a stack-local JointGraph (see - // decoder.cpp), which is what keeps concurrent decode reentrant. Owned here; - // freed by ~HostJoint. + // The whole joint runs as one ggml graph (build_joint_graph), so + // every weight is resident fp32: enc projection (g_enc_w/g_enc_b), + // pred projection (g_pred_w/g_pred_b), out projection (gw_w/gw_b). + // Built once at load, only read after — safe to share across every + // context; the mutable per-decode state lives in a stack-local + // JointGraph (reentrant). Owned here, freed by dtor. ggml_context * w_ctx = nullptr; ggml_backend_t w_backend = nullptr; // alloc-only; never graph_compute'd ggml_backend_buffer_t w_buf = nullptr; @@ -159,10 +142,7 @@ struct HostJoint { // CTC head mirror. NeMo's `decoder.decoder_layers.0` is a 1×1 Conv1d // projecting d_enc -> n_classes (= vocab + 1 blank). Stored row-major -// [n_classes, d_enc] so per-frame logits = W @ enc_t + b. The host -// decode path runs entirely on fp32 (same `read_tensor_to_f32` as the -// predictor/joint mirrors), so any of the GET_LIN-allowed source -// types are safe. +// [n_classes, d_enc] so per-frame logits = W @ enc_t + b. struct HostCtcHead { int n_classes = 0; // == vocab + 1 int blank_id = 0; // NeMo convention: blank lives at n_classes - 1 @@ -177,11 +157,10 @@ struct HostCtcHead { enum class HostHeadKind { TDT, RNNT, CTC }; // Bundle of everything the decoder needs at run time. Built once at -// load() time and stored on ParakeetModel; const-after-construction -// and shared across every context derived from the model. For -// head_kind=CTC the predictor + joint mirrors stay empty; for -// head_kind=RNNT the tdt_durations vector is empty and only blank/non-blank -// symbol-cap logic from `tdt_max_symbols` is reused. +// load() and stored on ParakeetModel; const-after-construction, shared +// across every context. For CTC the predictor + joint mirrors stay +// empty; for RNNT tdt_durations is empty and only the tdt_max_symbols +// cap is reused. struct HostDecoderWeights { HostHeadKind head_kind = HostHeadKind::TDT; HostPredictor predictor; // empty for CTC @@ -200,15 +179,10 @@ struct HostDecoderWeights { transcribe_status build_host_decoder_weights(const ParakeetModel & model, HostDecoderWeights & out); -// --------------------------------------------------------------------------- -// LSTM hidden state -// --------------------------------------------------------------------------- -// -// One (h, c) pair per LSTM layer. The decoder owns one "current" state -// (committed) and one "scratch" state (computed each iteration; only -// committed when a non-blank token is emitted). Both are sized once -// in reset() and reused across decode steps. - +// One (h, c) pair per LSTM layer. The decoder owns one "current" +// (committed) state and one "scratch" state (computed each iteration, +// committed only when a non-blank token is emitted). Sized once in +// reset() and reused across decode steps. struct LstmState { // h[layer] and c[layer] are each `pred_hidden` floats. std::vector> h; @@ -217,18 +191,11 @@ struct LstmState { void reset(int n_layers, int pred_hidden); }; -// --------------------------------------------------------------------------- -// Decode result -// --------------------------------------------------------------------------- -// // One emitted (non-blank) TDT token. `step_at_emit` is the encoder -// frame index when the token was emitted; `duration_frames` is the -// encoder-frame jump the joint network selected from the durations -// table. The decoder driver in Parakeet::run multiplies these by the -// model's frame-to-seconds ratio (subsampling_factor * hop_length / -// sample_rate) to convert to milliseconds for the public result -// accessors. - +// frame index at emission; `duration_frames` is the encoder-frame jump +// the joint selected from the durations table. The driver converts these +// to ms via the model's frame-to-seconds ratio (subsampling_factor * +// hop_length / sample_rate) for the public result accessors. struct TdtToken { int id = 0; float p = 0.0f; // entropy-based confidence @@ -239,29 +206,18 @@ struct TdtToken { // Run TDT greedy decode end-to-end against an encoder output buffer. // // Inputs: -// w - the decoder weights mirror -// enc_out - host pointer to T_enc * d_enc fp32 floats. Layout -// is fast-to-slow [d_enc, T_enc] = row-major [T_enc, d_enc]: -// frame `t` lives at `enc_out + t * d_enc`. Matches -// the byte layout of a ggml encoder output tensor with -// ne=[d_enc, T_enc, 1, 1] read via -// ggml_backend_tensor_get. -// T_enc - number of encoder frames -// d_enc - encoder d_model (must equal w.joint.d_enc) -// +// w - the decoder weights mirror +// enc_out - host pointer to T_enc * d_enc fp32 floats, fast-to-slow +// [d_enc, T_enc] = row-major [T_enc, d_enc]: frame t at +// enc_out + t * d_enc (the byte layout of a ggml tensor +// ne=[d_enc, T_enc, 1, 1] read via tensor_get). +// T_enc - number of encoder frames +// d_enc - encoder d_model (must equal w.joint.d_enc) // Outputs: -// out_tokens - appended with one TdtToken per non-blank emission +// out_tokens - appended with one TdtToken per non-blank emission // -// Side effects: -// When TRANSCRIBE_DUMP_DIR is set, dumps a small set of decoder -// intermediates (first-step embed, first-step LSTM h/c, first-step -// joint logits) for the bring-up comparison harness. Production -// builds (env var unset) pay zero cost. -// -// Returns TRANSCRIBE_OK on success. The decoder cannot fail at run -// time except via invalid arguments — input validation (T_enc > 0, -// d_enc matches the joint mirror, etc.) is enforced by the family -// driver before this is called. +// Returns TRANSCRIBE_OK on success; cannot fail except via invalid args +// (validated by the family driver before this is called). transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, const float * enc_out, int T_enc, @@ -270,15 +226,10 @@ transcribe_status decode_tdt_greedy(const HostDecoderWeights & w, std::vector & out_tokens); // Run RNNT greedy decode end-to-end. Same predictor + joint code as TDT, -// but the joint emits `vocab+1` logits (no duration extras) and the -// step rule is "blank → advance one frame, non-blank → emit + stay, -// capped by tdt_max_symbols". Per-emit `duration_frames` is fixed at 1 -// — the public token timestamps approximate the reference's -// (encoder-frame indexed) emission. -// -// Same input/output contract as decode_tdt_greedy. Reuses the same -// TdtToken result type so the model.cpp result-builder can stay -// head-agnostic. +// but the joint emits `vocab+1` logits (no duration extras) and the step +// rule is "blank → advance one frame, non-blank → emit + stay, capped by +// tdt_max_symbols". Per-emit duration_frames is fixed at 1. Same I/O +// contract and TdtToken result type as decode_tdt_greedy. transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, const float * enc_out, int T_enc, @@ -286,19 +237,13 @@ transcribe_status decode_rnnt_greedy(const HostDecoderWeights & w, int n_threads, std::vector & out_tokens); -// Streaming variant of RNN-T greedy decode. Consumes T_enc_new -// encoder frames (the chunk just produced by the streaming encoder -// graph) and APPENDS the emitted tokens to out_tokens. LSTM state and -// previous-token id are carried across calls via state_io and -// last_token_io. frame_offset is the absolute encoder-frame index of -// the first frame in this chunk (so emitted tokens' -// step_at_emit lands in stream-wide coordinates). -// -// state_io must have been reset to a fresh "start of sequence" state -// at stream_begin (LstmState::reset for n_layers, pred_hidden, plus -// last_token_io = -1 for the SOS embedding zero). -// -// No timing log is printed (per-chunk noise on the CLI). +// Streaming variant of RNN-T greedy decode. Consumes T_enc_new encoder +// frames (the chunk just produced) and APPENDS emitted tokens to +// out_tokens. LSTM state and previous-token id carry across calls via +// state_io / last_token_io; frame_offset is the absolute encoder-frame +// index of this chunk's first frame (so step_at_emit lands in +// stream-wide coordinates). state_io must have been reset to a fresh +// start-of-sequence state at stream_begin (last_token_io = -1). transcribe_status decode_rnnt_greedy_streaming( const HostDecoderWeights & w, const float * enc_out, @@ -311,14 +256,10 @@ transcribe_status decode_rnnt_greedy_streaming( std::vector & out_tokens); // Run CTC greedy decode end-to-end. Per-frame: logits = W @ enc[t] + b, -// argmax. The collapse rule is "drop adjacent duplicates, then drop -// blanks" — standard CTC greedy (see e.g. NeMo's GreedyCTCInfer). -// `step_at_emit` is the encoder-frame index of the (post-collapse) -// emission; `duration_frames` is set to 1 for compatibility with the -// public token timestamps. -// -// Same TdtToken result type as the transducer paths so model.cpp can -// build the public result hierarchy uniformly. +// argmax; collapse rule "drop adjacent duplicates, then drop blanks" +// (standard CTC greedy, cf. NeMo GreedyCTCInfer). step_at_emit is the +// post-collapse encoder-frame index; duration_frames = 1. Same TdtToken +// result type as the transducer paths. transcribe_status decode_ctc_greedy(const HostDecoderWeights & w, const float * enc_out, int T_enc, diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index e6897cca..4d680b24 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -1,44 +1,24 @@ // arch/parakeet/encoder.cpp - Parakeet Conformer encoder graph builder. // -// Thin glue over the shared Conformer helpers in src/conformer/. The -// per-op helpers (layer_norm, f32-friendly conv helpers, rel_shift, -// conv_module, rel_pos_mhsa, build_conformer_block, build_pre_encode) -// now live in transcribe::conformer — see src/conformer/conformer.{h,cpp} -// for the full implementations and the Metal backstory on the -// conv_*_f32 helpers. -// -// Parakeet-specific pieces kept in this file: -// - detect_direct_dw (different policy than Cohere) -// - to_view() helpers that project ParakeetPreEncode / ParakeetBlock -// onto the shared nullable-pointer views. Parakeet's conformer -// blocks ship without biases on the FFN/attention/conv-module -// linear layers, so those slots in BlockView stay nullptr — but -// pre_encode has real conv biases (conv0_b..conv6_b, out_b) and -// those ARE copied through by to_view(ParakeetPreEncode). +// Thin glue over the shared Conformer helpers in transcribe::conformer +// (src/conformer/conformer.{h,cpp}). Parakeet-specific pieces kept here: +// - detect_direct_dw (different policy than Cohere). +// - to_view() helpers projecting ParakeetPreEncode / ParakeetBlock onto +// the shared nullable-pointer views (block linear layers ship without +// biases; pre_encode convs carry biases that ARE copied through). // - build_encoder_graph orchestration — every block goes through -// conf::build_conformer_block; block 0 passes a BlockObserver -// that names intermediates and stashes them in the EncoderDumps -// struct so the 3b-3e bring-up harness can diff every sub-step +// conf::build_conformer_block; block 0 passes a dump observer. // // Layout cheat sheet: -// -// MelFrontend output : row-major [n_mels=128, n_frames=T_mel] f32 -// ggml mel_in tensor : ne=[T_mel, 128, 1, 1] f32 (matches the -// row-major buffer byte-for-byte) -// conv input form : [W=F, H=T, IC, N] per ggml_conv_2d. The -// catalog kernels are stored in PyTorch OIHW -// order which translates to ggml ne -// [KW, KH, IC, OC] where KW is aligned with -// the FREQ axis and KH with the TIME axis -// (NHWC convention: H=T, W=F). The 3x3 conv -// kernel is not spatially symmetric so the -// data must be presented with W=F, H=T. -// -// After 3 stride-2 convs (kernel=3, padding=1), W and H shrink by -// floor((d + 2 - 3)/2) + 1 per layer: -// F (ggml W) = 128 -> 64 -> 32 -> 16 -// T (ggml H) = 1101 -> 551 -> 276 -> 138 (jfk.wav, 11 s) -// So the post-conv tensor has ne = [F'=16, T_enc=138, channels=256, 1]. +// MelFrontend output : row-major [n_mels, T_mel] f32 +// ggml mel_in : ne=[T_mel, n_mels, 1, 1] f32 (byte-identical) +// conv input : [W=F, H=T, IC, N] per ggml_conv_2d; OIHW catalog +// kernels map to ggml ne [KW, KH, IC, OC] with KW +// on the FREQ axis and KH on the TIME axis (the 3x3 +// kernel isn't symmetric, so W=F, H=T is required). +// After 3 stride-2 k=3 pad=1 convs, W and H shrink floor((d+2-3)/2)+1 +// per layer (n_mels=128: F 128→64→32→16); post-conv ne = +// [F', T_enc, channels, 1]. #include "encoder.h" @@ -63,38 +43,25 @@ namespace { namespace conf = transcribe::conformer; -// ----- Per-family depthwise dispatch policy ----------------------- +// ----- Per-family depthwise dispatch policy ----- // -// Parakeet uses the direct conv_2d_dw path on every backend for the -// in-block depthwise conv (Vulkan: native kernel; Metal/CPU: the -// direct path is faster than the 31x im2col expansion). The -// pre_encode site historically hardcoded the im2col f32 helper -// regardless of the in-block flag, so direct_dw_in_pre_encode stays -// false here to preserve that behavior exactly. Env vars map onto -// both sites because Parakeet has no history of splitting them. +// In-block depthwise conv uses the direct conv_2d_dw path on every +// backend (faster than the 31x im2col expansion). Env vars override +// both sites. bool detect_direct_dw_in_block(const char * backend) { const char * env = std::getenv("TRANSCRIBE_CONV_DIRECT_DW"); if (env != nullptr) return true; // user override env = std::getenv("TRANSCRIBE_CONV_NO_DIRECT_DW"); if (env != nullptr) return false; // user override - // Vulkan and CPU have native CONV_2D_DW; Metal does too via the - // direct path in this version. CPU also benefits from direct dw - // (avoids the 31x im2col expansion). (void)backend; return true; } -// Pre-encode (subsampling) depthwise dispatch. Same direct path as the -// in-block site EXCEPT on Metal, where ggml_conv_2d_dw_direct's CONV_2D_DW -// op is unsupported in this vendored ggml (see conv_2d_dw_f32's comment in -// conformer.cpp) so the im2col helper must be used. CPU and Vulkan have a -// native CONV_2D_DW: the direct path avoids both the ~31x im2col expansion -// AND the degenerate per-channel [1 x K] batched matmul the im2col path -// emits (m=1, batch=channels) — a large, GPU-/CPU-idle pre-encode cost -// (~40-50ms here). Historically this was hardcoded to the im2col helper to -// preserve Metal behavior; making it backend-aware keeps Metal correct -// while letting CPU/Vulkan take the fast path. Env vars match the in-block -// overrides. +// Pre-encode depthwise dispatch: direct path EXCEPT on Metal, where +// CONV_2D_DW is unsupported in this vendored ggml (see conv_2d_dw_f32 in +// conformer.cpp) so the im2col helper is used. CPU/Vulkan have native +// CONV_2D_DW; the direct path avoids the ~31x im2col expansion and the +// degenerate per-channel [1 x K] matmul (~40-50ms here). bool detect_direct_dw_in_pre_encode(const char * backend) { if (std::getenv("TRANSCRIBE_CONV_DIRECT_DW") != nullptr) return true; if (std::getenv("TRANSCRIBE_CONV_NO_DIRECT_DW") != nullptr) return false; @@ -129,9 +96,8 @@ conf::BlockView to_view(const ParakeetBlock & b) { v.ff1_lin2_w = b.ff1_lin2_w; v.ff1_lin2_b = b.ff1_lin2_b; - // Self-attention. Q/K/V/out can carry biases when use_bias=true - // (1.1B / rnnt / ctc variants); v2/v3 ship without. linear_pos is - // bias-free in NeMo regardless. + // Self-attention. Q/K/V/out biases only when use_bias=true; linear_pos + // is bias-free in NeMo regardless. v.norm_attn_w = b.norm_attn_w; v.norm_attn_b = b.norm_attn_b; v.attn_q_w = b.attn_q_w; v.attn_q_b = b.attn_q_b; @@ -153,12 +119,9 @@ conf::BlockView to_view(const ParakeetBlock & b) { v.conv_pw2_b = b.conv_pw2_b; v.conv_bn_fused_scale = b.conv_bn_fused_scale; v.conv_bn_fused_bias = b.conv_bn_fused_bias; - // For the LayerNorm path (streaming variants) the GGUF stores the - // affine scale/bias under the same conv.bn.weight/bias tensor names - // (NeMo keeps the Python attribute named `batch_norm` even when the - // module is swapped to LayerNorm). Point conv_ln_* at the raw - // tensors; the conformer block only reads them when - // params.conv_norm_type == LayerNorm. + // LayerNorm path (streaming variants): the GGUF stores the affine + // scale/bias under the same conv.bn.weight/bias names. The conformer + // block reads these only when conv_norm_type == LayerNorm. v.conv_ln_w = b.conv_bn_w; v.conv_ln_b = b.conv_bn_b; @@ -175,20 +138,13 @@ conf::BlockView to_view(const ParakeetBlock & b) { return v; } -// ----- Block 0 observer ------------------------------------------- +// ----- Block 0 observer ----- // -// Drives the 3b-3e bring-up harness. For each of the five canonical -// points build_conformer_block fires on, we name the tensor -// `enc.block.0.` (matching the old hand-built block's naming -// exactly, so existing graph-inspection workflows keep working) and -// stash the pointer in the right EncoderDumps slot so Parakeet::run's -// try_dump loop finds it. -// -// NOTE: the dump FILE names in model.cpp's try_dump table -// (`enc.block.0.ff1`, etc.) are independent of the ggml names here -// (`enc.block.0.after_ff1`, etc.). That pre-dates the observer — -// both names existed on the hand-built block as well. Preserved -// as-is to avoid any behavioral change during the refactor. +// For each canonical point build_conformer_block fires on, names the +// tensor `enc.block.0.` and stashes it in the matching EncoderDumps +// slot for Parakeet::run's try_dump loop. The dump FILE names in +// model.cpp (enc.block.0.ff1, ...) are independent of these ggml names +// (enc.block.0.after_ff1, ...). struct Block0DumpSink { EncoderDumps * dumps; }; @@ -222,11 +178,10 @@ void parakeet_block0_observer_cb(void * user, } } -// Generalized sub-block observer for blocks 1..N-1, gated by -// TRANSCRIBE_DUMP_SUB_BLOCKS=. Names tensors -// `enc.block..{ff1,attn,conv,ff2,out}` (mapping the helper's "after_*" -// tags to short names matching the reference dumper). The "out" tag -// here corresponds to the post-norm-out tensor — i.e. block..out. +// Sub-block observer for blocks 1..N-1, gated by +// TRANSCRIBE_DUMP_SUB_BLOCKS. Names tensors +// `enc.block..{ff1,attn,conv,ff2}` (the helper's "out" tag is already +// covered by the mid/last/all-block paths). struct SubBlockSink { int block_idx; EncoderDumps * dumps; @@ -245,9 +200,7 @@ void parakeet_subblock_observer_cb(void * user, else if (std::strcmp(tag, "after_conv") == 0) short_tag = "conv"; else if (std::strcmp(tag, "after_ff2") == 0) short_tag = "ff2"; else if (std::strcmp(tag, "out") == 0) { - // Skip "out" — block..out is already covered by the - // mid_block_out / last_block_out / all_block_outs paths. - return; + return; // block..out covered elsewhere } char name_buf[64]; @@ -297,24 +250,14 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, { if (n_batch < 1) n_batch = 1; const bool var_len_masks = batch_var_len && n_batch > 1; - // Per-family dispatch policy. Parakeet keeps the historical split - // where the in-block depthwise uses the direct path on every - // backend but the pre_encode depthwise uses the f32-friendly - // im2col helper unconditionally. See detect_direct_dw_in_block - // above and the ConvPolicy comment in conformer.h. 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 = detect_direct_dw_in_pre_encode(backend_name); - // Cache-aware streaming variants (NeMo `causal_downsampling=true`) - // use CausalConv2D for the pre-encode subsample stack (left=k-1, - // right=stride-1 pad on both spatial axes). We infer this from the - // attention style: ChunkedLimited (cache-aware) implies causal - // subsample; Regular (offline) and ChunkedLimitedWithRc (buffered - // streaming) both use the standard non-causal symmetric pad. The - // pre-encode kernel (k=3) and the conformer conv-module kernel - // (k=9) are independent — hp.enc_conv_context_{left,right} are - // for the latter and don't say anything about the former. + // Cache-aware streaming (NeMo causal_downsampling=true) uses + // CausalConv2D for the pre-encode subsample (left=k-1, right=stride-1). + // Inferred from the attention style — only ChunkedLimited is causal. + // Independent of the conformer conv-module's conv_context. policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); @@ -329,12 +272,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } - // We support subsampling_factor=8 (every published Parakeet variant) - // and any n_mels divisible by the subsampling factor on the freq axis - // (3 stride-2 convs, kernel=3, padding=1). All current variants ship - // n_mels=80 or 128; both pass cleanly. If a future variant ships a - // factor != 8, build_pre_encode would need a structural rework, so - // we fail loudly here. + // Only subsampling_factor=8 (every published variant) is implemented; + // a different factor would need a build_pre_encode rework. if (hp.enc_subsampling_factor != 8) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet encoder: unsupported subsampling_factor=%d " @@ -350,12 +289,9 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } - // Mel input handle. ne=[T_mel, n_mels, 1, n_batch] matches the C++ - // MelFrontend's row-major [n_mels, n_frames] storage byte for byte - // (see the layout cheat sheet at the top of the file); utterance b is - // the contiguous [n_mel_frames * n_mels] slab at offset - // b * n_mel_frames * n_mels. For n_batch == 1 this is identical to the - // prior 2-D [T_mel, n_mels] tensor. + // Mel input handle. ne=[T_mel, n_mels, 1, n_batch] matches the + // MelFrontend's row-major [n_mels, n_frames] byte for byte; utterance + // b is the contiguous slab at offset b * n_mel_frames * n_mels. eb.mel_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_mel_frames, hp.fe_num_mels, 1, n_batch); if (eb.mel_in == nullptr) { @@ -365,19 +301,13 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } ggml_set_name(eb.mel_in, "mel.in"); ggml_set_input(eb.mel_in); - // Debug-only: pin the mel_in buffer so the scheduler doesn't reuse - // its slot for a later activation before dump_tensor() reads it - // back. Without this, streaming dumps of enc.mel.in show garbage - // because n_mel_frames is small and the slot is the first to be - // freed. No-op in normal builds. + // Pin mel_in so the scheduler doesn't reuse its slot before the dump + // pass reads it (no-op in normal builds). transcribe::debug::mark_tensor_for_dump(eb.mel_in); - // Build the pre_encode subgraph. For variable-length offline batching we - // enable masked subsampling (zero each conv intermediate's padded time - // region so an utterance's boundary frame matches a standalone run). Only - // the non-causal pre-encode is masked here — the causal/streaming stem - // has a different per-stage length formula and its boundary frame does - // not flip decoded tokens in practice. + // Build the pre_encode subgraph. Variable-length offline batching + // enables masked subsampling (zero each conv intermediate's padded + // time region). Only the non-causal pre-encode is masked. conf::PreEncodeValidMasks pe_masks; const bool mask_pre_encode = var_len_masks && !policy.causal_pre_encode; ggml_tensor * x = conf::build_pre_encode(ctx, to_view(w.pre_encode), @@ -397,67 +327,36 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.pre_encode_out = x; transcribe::debug::mark_tensor_for_dump(x); - // ----- xscaling --------------------------------------------------- - // - // NeMo's RelPositionalEncoding.forward applies x = x * sqrt(d_model) - // when its `xscaling` flag is true. This multiplication happens AFTER - // pre_encode and BEFORE the first conformer block — i.e. it - // operates on the running residual, not just on the pos_emb input. - // ctc-0.6b/1.1b, rnnt-0.6b/1.1b, and unified-en-0.6b ship with - // xscaling=True; v2/v3 and the TDT/TDT_CTC variants ship with - // xscaling=False. + // ----- xscaling ----- // - // The same shape comes out (the scale factor is applied uniformly - // along d_model). LayerNorms in every sub-block normalize the - // scale away when computing FF/attn/conv outputs, so the only - // place the scale matters is the residual carried between - // sub-blocks and the running residual into norm_out at the end of - // each block. Without this multiplication on a True-flagged model - // the residual is sqrt(d_model)x smaller, the corrections are the - // same magnitude (because they go through LN first), and the - // residual+correction mix is structurally different — the - // accumulated drift past 24 blocks is enough to produce empty - // transcripts on unified-en-0.6b and ~0.04% WER drift on ctc-0.6b - // / rnnt-0.6b before this fix landed. + // NeMo's RelPositionalEncoding applies x = x * sqrt(d_model) when + // xscaling=True, on the running residual after pre_encode and before + // block 0 (ctc/rnnt/unified-en are True; v2/v3/tdt are False). + // Sub-block LayerNorms normalize the scale away, so it only matters + // for the residual mix. Omitting it on a True model accumulates drift + // past 24 blocks — empty transcripts on unified-en, ~0.04% WER drift + // on ctc/rnnt. if (hp.enc_xscaling) { const float scale = std::sqrt(static_cast(hp.enc_d_model)); x = ggml_scale(ctx, x, scale); x = conf::named(x, "enc.pre_encode.xscaled"); } - // ----- Conformer blocks ----------------------------------------- + // ----- Conformer blocks ----- // - // Every block goes through conf::build_conformer_block — there is - // no hand-built path anymore. Block 0 passes a BlockObserver that - // names intermediates and stashes them in eb.dumps so the 3b-3e - // bring-up harness can diff every sub-step. Blocks 12 and 23 get - // their final-LN outputs named after return; the loop is plain - // otherwise so the dump dir doesn't fill with all 24 outputs. - + // Every block goes through conf::build_conformer_block. Block 0 passes + // a dump observer; mid/last blocks get their outputs named after + // return (the loop is otherwise plain so the dump dir doesn't fill). if (!w.blocks.empty()) { - // After pre_encode, x ne = [d_model, T_enc, 1, 1]. T_enc is - // needed by the positional embedding (whose pos_len = - // 2*T_enc - 1) and by the encoder accuracy harness for shape - // sanity. + // After pre_encode, x ne = [d_model, T_enc, 1, 1]. const int64_t T_enc = x->ne[1]; - // Create the pos_emb input tensor here, once we know T_enc. - // The driver fills it via ggml_backend_tensor_set after the - // compute buffer is allocated. ne = [d_model, pos_len, 1, 1] - // — slow-to-fast shape (numpy) is (pos_len, d_model), - // matching the reference `enc.pos_emb` dump. - // - // Local attention (Regular style with both att_context_{left,right} - // >= 0) shortens pos_emb to (left+right+1) so the buffer covers - // exactly the attended relative-position range — matches NeMo's - // LocalAttRelPositionalEncoding. ChunkedLimited style keeps pos_emb - // at the full 2T-1 and uses a separate attention-position mask - // (built below). - // ChunkedLimited (2-tuple cache-aware) always engages the mask. - // ChunkedLimitedWithRc (3-tuple buffered streaming) only engages - // when the caller passes a non-null buf_mask override — offline - // load of a unified-en GGUF runs the encoder with full attention - // even though the style declares chunked_limited_with_rc. + // pos_emb input, ne = [d_model, pos_len, 1, 1] (numpy + // (pos_len, d_model)), filled by the driver. Local attention + // (Regular, both contexts >= 0) shortens pos_len to (left+right+1) + // (NeMo LocalAttRelPositionalEncoding); ChunkedLimited keeps the + // full 2T-1 and uses a separate mask. ChunkedLimitedWithRc engages + // the mask only when buf_mask is non-null (offline runs full attn). const bool is_chunked = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited) || @@ -552,9 +451,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ? conf::BlockParams::ConvNormType::LayerNorm : conf::BlockParams::ConvNormType::BatchNorm; - // Block 0: run through the shared helper with the dump - // observer. The observer writes into eb.dumps; no sub-step - // references are lost vs. the old hand-built path. + // Block 0: shared helper + the dump observer. { Block0DumpSink sink { &eb.dumps }; conf::BlockObserver obs {}; @@ -565,16 +462,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, bparams, &obs); } - // Blocks 1..N-1: no observer. Mark the mid- and last-layer - // outputs for dump so the harness can spot-check them; every - // other block stays anonymous. Mid + last indices scale with - // n_layers so the dump points match the per-variant oracle - // (0/n_half/n-1). For 24-layer v2/v3 that's 0/12/23; for - // 42-layer 1.1B it's 0/21/41; for 17-layer tdt_ctc-110m it's - // 0/8/16. We only mark for dump here — the actual file name - // is synthesized in model.cpp from the saved layer index, - // because conf::named's trailing "enc.final" rename would - // otherwise clobber the last-block name in place. + // Blocks 1..N-1. Mark the mid (n/2) and last (n-1) outputs for + // dump (0/12/23 on 24-layer v2/v3); the file name is synthesized + // in model.cpp from the saved index because conf::named's trailing + // "enc.final" rename would otherwise clobber the last-block name. const int n_layers = static_cast(w.blocks.size()); const int mid_layer = n_layers / 2; const int last_layer = n_layers - 1; @@ -582,29 +473,20 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.last_block_idx = last_layer; eb.dumps.all_block_outs.assign(n_layers, nullptr); eb.dumps.all_block_outs[0] = eb.dumps.block0_out; - // Per-block dump opt-in for the layer-by-layer divergence - // bisect. The reference dumper supports `--blocks 0,1,...,N-1`; - // setting TRANSCRIBE_DUMP_ALL_BLOCKS=1 makes the C++ side - // dump every block to enable a 1:1 frame-by-frame compare. + // Per-block dump opt-in for the divergence bisect + // (TRANSCRIBE_DUMP_ALL_BLOCKS). const bool dump_all_blocks = std::getenv("TRANSCRIBE_DUMP_ALL_BLOCKS") != nullptr; if (dump_all_blocks && eb.dumps.block0_out != nullptr) { transcribe::debug::mark_tensor_for_dump(eb.dumps.block0_out); } - // Sub-block observation gate. TRANSCRIBE_DUMP_SUB_BLOCKS= - // "12,22,23" installs the sub-block observer on each listed - // index >= 1 (block 0 already has its dedicated observer above). - // Stored in a local set for O(1) lookup in the loop. + // Sub-block observation gate (TRANSCRIBE_DUMP_SUB_BLOCKS). Sinks + // are kept alive for the build (the observer captures their addr). const std::vector sub_blocks = parse_sub_block_env(n_layers); - // Persistent sinks — one per requested block, kept alive for the - // duration of the encoder build. The observer pointer captures - // their address. std::vector sub_sinks; sub_sinks.reserve(sub_blocks.size()); for (int idx : sub_blocks) sub_sinks.push_back(SubBlockSink{idx, &eb.dumps}); for (int i = 1; i < n_layers; ++i) { - // If this block is requested for sub-block dumping, install - // the generalized observer; otherwise run plain. const auto it = std::lower_bound(sub_blocks.begin(), sub_blocks.end(), i); const bool with_obs = @@ -637,48 +519,31 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // Final encoder output. With all 24 blocks wired, this is now - // the true encoder exit point (= block 23's norm_out output) and - // the comparator's `enc.final` row will compare against the - // reference's true encoder output. + // Final encoder output (= last block's norm_out). eb.dumps.final_out = x; x = conf::named(x, "enc.final"); transcribe::debug::mark_tensor_for_dump(x); - // Force final_out to remain an output so the host can read it back - // separately from the prompted output below. Without this the - // scheduler may free the buffer once eb.out (= prompted) is - // realised and the host-side "dec.enc_out" dump would race the - // reuse. ggml_set_output also lets ggml_backend_tensor_get attach - // to the materialised slot. + // Keep final_out as an output so the host can read it back separately + // from the prompted output (else the scheduler may free it once eb.out + // is realised, racing the "dec.enc_out" dump). if (hp.has_prompt) { ggml_set_output(x); } - // Prompt MLP (multilingual variants only — today - // nemotron-3.5-asr-streaming-0.6b). Mirrors NeMo's + // Prompt MLP (multilingual variants). Mirrors NeMo's // EncDecRNNTBPEModelWithPrompt.forward(): - // - // x_cat = concat(enc[d_model], one_hot(prompt_id)[num_prompts]) - // h = relu(W0 @ x_cat + b0) // W0: [prompt_hidden, d_model+P] - // y = W2 @ h + b2 // W2: [d_model, prompt_hidden] - // enc_out := y - // - // The one-hot vector is replicated across the T_enc axis on the - // host (small buffer, num_prompts × T_enc × n_batch floats per - // call) so the in-graph step is plain concat + two matmuls + a - // ReLU. After the MLP the post-prompt output replaces eb.out so - // the decoder consumes the prompt-conditioned tensor unchanged - // from the non-prompt code path. + // x_cat = concat(enc[d_model], one_hot(prompt_id)[num_prompts]) + // h = relu(W0 @ x_cat + b0) // W0: [prompt_hidden, d_model+P] + // y = W2 @ h + b2 // W2: [d_model, prompt_hidden] + // The one-hot is replicated across T_enc on the host so the in-graph + // step is concat + two matmuls + ReLU; y replaces eb.out. if (hp.has_prompt) { const int T_enc_val = static_cast(x->ne[1]); const int P = hp.prompt_num_prompts; const int B = n_batch; - // Input: per-utterance one-hot replicated across T frames. - // ne = [num_prompts, T_enc, n_batch, 1] — batch rides ne[2] to - // match the Conformer block output layout - // ([d_model, T_enc, B, 1]; see conformer.cpp reshape_3d). Memory - // layout is identical to [P, T_enc, 1, B] for the host filler. + // Per-utterance one-hot, ne = [num_prompts, T_enc, n_batch, 1] + // (batch on ne[2] to match the block output layout). ggml_tensor * one_hot = ggml_new_tensor_4d( ctx, GGML_TYPE_F32, P, T_enc_val, B, 1); if (one_hot == nullptr) { @@ -691,16 +556,10 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, ggml_set_input(one_hot); eb.prompt_one_hot_in = one_hot; - // Concat along the feature axis (ne[0]) → - // [d_model + P, T_enc, 1, B] ggml_tensor * cat = ggml_concat(ctx, x, one_hot, /*dim=*/0); - - // Linear mlp.0: W0 @ cat + b0 → [prompt_hidden, T_enc, 1, B] ggml_tensor * h = ggml_mul_mat(ctx, w.prompt.mlp0_w, cat); h = ggml_add(ctx, h, w.prompt.mlp0_b); h = ggml_relu(ctx, h); - - // Linear mlp.2: W2 @ h + b2 → [d_model, T_enc, 1, B] ggml_tensor * y = ggml_mul_mat(ctx, w.prompt.mlp2_w, h); y = ggml_add(ctx, y, w.prompt.mlp2_b); @@ -714,13 +573,8 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.out = x; ggml_set_output(eb.out); - // Build the forward cgraph. ggml_new_graph_custom allocates the - // cgraph out of the same compute_ctx, so it has to be sized for - // it. The default GGML_DEFAULT_GRAPH_SIZE (2048 nodes) is too - // small for the full encoder: each conformer block contributes - // ~90 ops (FF + attn + conv + FF + LN, including all the - // permute/cont overhead) and 24 blocks pushes the count past - // 2200. Use 8192 to leave headroom. + // 8192-node cgraph: the default 2048 is too small (each conformer + // block contributes ~90 ops, 24 blocks > 2200). eb.graph = ggml_new_graph_custom(ctx, /*size=*/8192, /*grads=*/false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -728,11 +582,9 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } ggml_build_forward_expand(eb.graph, eb.out); - // When the prompt MLP is present, `eb.out` already feeds into the - // prompted output, but `eb.dumps.final_out` is an unconnected node - // from the cgraph's POV. Adding it as a forward root keeps the - // un-prompted tensor materialised so the host can read it back for - // the "dec.enc_out" dump. + // With the prompt MLP present, final_out is otherwise an unconnected + // node; add it as a forward root so it stays materialised for the + // "dec.enc_out" dump. if (hp.has_prompt && eb.dumps.final_out != nullptr) { ggml_build_forward_expand(eb.graph, eb.dumps.final_out); } @@ -744,27 +596,22 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Streaming encoder graph (cache-aware, per-chunk feed) // --------------------------------------------------------------------------- // -// Mirrors build_encoder_graph but: -// - Operates on a mel chunk (n_mel_chunk_frames), not the whole audio. -// - Slices off `drop_extra_pre_encoded` post-subsample frames. -// - Threads per-layer cache_last_channel + cache_last_time tensors -// through BlockParams, and emits per-layer cache outputs into -// fresh compute_ctx tensors that the driver reads back. -// - Sizes pos_emb and chunked_mask for T_virtual = T_cache + T_q_new -// (the per-block attention runs on the cache+new union). +// Like build_encoder_graph but: operates on a mel chunk; slices off +// drop_extra_pre_encoded post-subsample frames; threads per-layer +// cache_last_channel/time through BlockParams and emits fresh cache +// outputs the driver reads back; sizes pos_emb and chunked_mask for +// T_virtual = T_cache + T_q_new. // // Shapes: -// eb.mel_in [n_mel_chunk_frames, n_mels, 1, 1] -// eb.pos_emb_in [d_model, 2*T_virtual - 1, 1, 1] -// eb.chunked_mask_in [T_virtual, T_virtual, 1, 1] -// eb.out [d_model, T_q_new, 1, 1] -// cache_io.channel_out[i] [d_model, T_cache, 1, 1] (fresh) -// cache_io.time_out[i] [k_minus_1, d_model, 1, 1] (fresh) +// eb.mel_in [n_mel_chunk_frames, n_mels, 1, 1] +// eb.pos_emb_in [d_model, 2*T_virtual - 1, 1, 1] +// eb.chunked_mask_in [T_virtual, T_virtual, 1, 1] +// eb.out [d_model, T_q_new, 1, 1] +// cache_io.channel_out[i] [d_model, T_cache, 1, 1] (fresh) +// cache_io.time_out[i] [k_minus_1, d_model, 1, 1] (fresh) // -// All cache input tensors (cache_io.channel_in / time_in) must come -// from the persistent ParakeetStreamingCaches buffer and have the -// shapes above. The graph DOES treat them as inputs (the scheduler -// walks back from cpy ops and eb.out to find them). +// cache_io.channel_in / time_in must come from the persistent cache +// buffer with the shapes above (the scheduler treats them as inputs). EncoderBuild build_encoder_graph_streaming( ggml_context * ctx, const ParakeetWeights & w, @@ -779,8 +626,7 @@ EncoderBuild build_encoder_graph_streaming( 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 = detect_direct_dw_in_pre_encode(backend_name); - // See the offline analog in build_encoder_graph: causal pre-encode - // is the cache-aware streaming convention only. + // Causal pre-encode is the cache-aware streaming convention only. policy.causal_pre_encode = (hp.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); @@ -865,20 +711,14 @@ EncoderBuild build_encoder_graph_streaming( const int64_t T_cache = cache_io.channel_in[0]->ne[1]; const int64_t T_virt = T_cache + T_q_new; - // Query slicing (always on): attention queries cover only the - // T_q_new new frames while K/V span the full virtual window, so - // pos_emb and the chunked mask are sized for that rectangular - // [T_virt, T_q_new] geometry. The host-side fills in - // emit_streaming_chunk read the sizes back off these tensors, so - // builder and driver stay in sync by construction. - - // Streaming KV cache: consume per-layer pre-projected K/V instead - // of recomputing them from the channel cache (K/V projections then - // run on T_q_new columns instead of T_virt). Disabled when the dump - // harness is active — the numerical validation compares last_channel - // snapshots against NeMo, and this mode neither reads nor writes - // them, so dumping falls back to the channel-recompute path (which - // produces bit-identical output and the last_channel the gate needs). + // Query slicing: attention queries cover only the T_q_new new frames + // while K/V span the full virtual window, so pos_emb and the mask are + // sized for the rectangular [T_virt, T_q_new] geometry. + + // Streaming KV cache: consume per-layer pre-projected K/V instead of + // recomputing from the channel cache. Disabled under the dump harness + // (which compares last_channel against NeMo); the channel-recompute + // path is bit-identical and produces the last_channel the gate needs. const bool kv_cache_mode = !transcribe::debug::enabled() && static_cast(cache_io.k_in.size()) == n_layers && @@ -906,33 +746,24 @@ EncoderBuild build_encoder_graph_streaming( ggml_set_name(eb.chunked_mask_in, "stream.attn.chunked_mask.in"); ggml_set_input(eb.chunked_mask_in); - // ----- Create the graph BEFORE allocating cache_out tensors, so - // the helpers can ggml_build_forward_expand their cpy nodes onto - // it. ----- + // Create the graph BEFORE the cache_out tensors so the helpers can + // forward-expand their cpy nodes onto it. eb.graph = ggml_new_graph_custom(ctx, /*size=*/8192, /*grads=*/false); if (eb.graph == nullptr) return eb; - // ----- BlockParams (per-block, but only streaming_*_in/_out and - // streaming_T_q_new change per layer) ----- conf::BlockParams bparams {}; bparams.d_model = hp.enc_d_model; bparams.n_head = hp.enc_n_heads; bparams.conv_kernel = hp.enc_conv_kernel; bparams.kv_type = kv_type; - // Same flash-attention opt-out as the offline builder. On CPU the - // flash kernel is dramatically slower than the manual mul_mat + - // soft_max path at streaming geometry (T_q ~14, T_k ~70): measured - // ~33 ms/layer vs ~1 ms on a Cortex-A55. The numerics gate - // (flash casts the rel-pos mask to F16, manual stays F32) is the - // same one the offline path documents. + // Flash-attention opt-out (TRANSCRIBE_NO_FLASH). On CPU at streaming + // geometry the flash kernel is dramatically slower (~33 ms/layer vs + // ~1 ms on a Cortex-A55). bparams.use_flash = (std::getenv("TRANSCRIBE_NO_FLASH") == nullptr); bparams.policy = policy; - // No att_context_left/right here: ChunkedLimited derives the attention - // band entirely from attn_chunked_mask (built host-side in - // emit_streaming_chunk). build_conformer_block only reads - // att_context_left/right on the Regular (is_local) path, which this - // builder never takes — leaving them at the BlockParams defaults keeps - // the live source of truth (the mask) singular. + // No att_context_left/right: ChunkedLimited derives the band entirely + // from attn_chunked_mask (build_conformer_block reads them only on the + // Regular path, never taken here). bparams.att_context_style = conf::BlockParams::AttContextStyle::ChunkedLimited; bparams.attn_chunked_mask = eb.chunked_mask_in; bparams.conv_context_left = hp.enc_conv_context_left; @@ -953,13 +784,9 @@ EncoderBuild build_encoder_graph_streaming( cache_io.v_out.assign(n_layers, nullptr); } for (int i = 0; i < n_layers; ++i) { - // Per-layer cache I/O. The "in" tensors are from the - // persistent cache buffer (passed in by the driver). The - // "out" tensors are fresh in compute_ctx; the helpers fill - // them via ggml_cpy and the driver reads them after compute. - // KV mode swaps the channel cache for pre-projected K/V (the - // channel_out slot stays null; the driver rotates whichever - // outs are present). + // Per-layer cache I/O: "in" from the persistent buffer, "out" + // fresh in compute_ctx (filled via ggml_cpy, read after compute). + // KV mode leaves channel_out null. ggml_tensor * cache_tm_out = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, k_minus_1, hp.enc_d_model); if (cache_tm_out == nullptr) return eb; diff --git a/src/arch/parakeet/encoder.h b/src/arch/parakeet/encoder.h index afdb05f1..2a061f6c 100644 --- a/src/arch/parakeet/encoder.h +++ b/src/arch/parakeet/encoder.h @@ -1,27 +1,9 @@ // arch/parakeet/encoder.h - Parakeet Conformer encoder graph builder. // -// Phase 4 step 3 of the encoder port. The .cpp file builds a -// ggml_cgraph that mirrors NeMo's Conformer encoder forward. The C++ -// encoder reads dims from ParakeetHParams and tensor pointers from -// ParakeetWeights; both are populated by the loader. -// -// Sub-stages of step 3: -// -// 3a: pre_encode subsampling stack only. Validates layout + conv2d -// wiring + linear projection + the compute-context lifetime -// against the reference enc.pre_encode.out on samples/jfk.wav. -// 3b: macaron FF1 on block 0. -// 3c: relative-position MHSA on block 0. -// 3d: conv module on block 0. -// 3e: FF2 + final norm on block 0. -// 3f: loop over 24 blocks, validate against block 1, 12, 23, final. -// -// Each sub-stage adds named dump points consumed by -// scripts/compare_tensors.py. The first-divergent-block-wins debug -// strategy is the entire dev loop. -// -// This header is INTERNAL to src/arch/parakeet/. The public C ABI -// only sees Parakeet::run, which builds and runs the graph below. +// The .cpp builds a ggml_cgraph mirroring NeMo's Conformer encoder +// forward, reading dims from ParakeetHParams and tensor pointers from +// ParakeetWeights (both populated by the loader). Internal to +// src/arch/parakeet/; the public C ABI only sees Parakeet::run. #pragma once @@ -40,52 +22,38 @@ namespace transcribe::parakeet { struct ParakeetHParams; struct ParakeetWeights; -// Result of building one pass of the encoder graph. The caller is -// responsible for: -// -// 1. allocating a backend buffer for the compute_ctx (typically via -// ggml_backend_alloc_ctx_tensors against the model's backend), -// 2. uploading the mel data into `mel_in` via ggml_backend_tensor_set, -// 3. running the cgraph via ggml_backend_graph_compute, -// 4. reading or dumping `out` via ggml_backend_tensor_get / the -// transcribe::debug dumper. +// Result of building one pass of the encoder graph. The caller +// allocates a backend buffer for compute_ctx, uploads mel into `mel_in`, +// runs the cgraph, then reads/dumps `out`. // -// Layout convention (see RESUME.md "Where we are" for the long form): -// - The reference uses [B, T, C] (channels-last). ggml ne is -// fast-to-slow, so the natural encoder activation lives at +// Layout convention: +// - The reference uses [B, T, C] (channels-last); ggml ne is +// fast-to-slow, so the natural encoder activation is // ne=[d_model, T, B=1, 1]. -// - For pre_encode (Conv2d), we treat T_mel as W (ggml ne[0]) and -// n_mels as H (ne[1]); this matches the row-major layout the -// C++ MelFrontend already produces, so the input tensor receives -// the mel buffer with no transpose. -// - After three stride-2 convs, the encoder activation enters the -// conformer blocks at ne=[d_model, T_enc, 1, 1] where T_enc = -// floor(T_mel / 8) (with the per-conv padding/kernel formula -// applied three times — see encoder.cpp). -// Named intermediate dump points the encoder graph builder -// publishes. The driver in Parakeet::run iterates this set after -// graph_compute and feeds each tensor to transcribe::debug::dump_tensor -// (gated on TRANSCRIBE_DUMP_DIR; zero-cost when unset). Any tensor -// stored here must also be passed to transcribe::debug::mark_tensor_for_dump() -// while building the graph; otherwise the scheduler may reuse its -// buffer before the post-compute dump pass reads it. Adding a new -// sub-stage means appending to this struct + populating/preserving it -// in build_encoder_graph; Parakeet::run doesn't need to know the names. +// - For pre_encode (Conv2d), T_mel is W (ne[0]) and n_mels is H +// (ne[1]); this matches the C++ MelFrontend's row-major layout, so +// the input tensor receives the mel buffer with no transpose. +// - After three stride-2 convs, the activation enters the conformer +// blocks at ne=[d_model, T_enc, 1, 1] where T_enc = floor(T_mel / 8) +// (per-conv padding/kernel formula applied three times). +// Named intermediate dump points the encoder graph builder publishes. +// Parakeet::run iterates this after graph_compute and feeds each tensor +// to transcribe::debug::dump_tensor (gated on TRANSCRIBE_DUMP_DIR). Any +// tensor stored here must also be passed to mark_tensor_for_dump() while +// building, or the scheduler may reuse its buffer before the dump pass. struct EncoderDumps { - // Pre-encode (3a). ne=[d_model, T_enc, 1, 1]. + // Pre-encode. ne=[d_model, T_enc, 1, 1]. ggml_tensor * pre_encode_out = nullptr; - // Block 0 sub-step outputs (3b-3e). Each ne=[d_model, T_enc, 1, 1]. + // Block 0 sub-step outputs. Each ne=[d_model, T_enc, 1, 1]. ggml_tensor * block0_after_ff1 = nullptr; ggml_tensor * block0_after_attn = nullptr; ggml_tensor * block0_after_conv = nullptr; ggml_tensor * block0_after_ff2 = nullptr; ggml_tensor * block0_out = nullptr; - // Spot-check blocks (3f). Mid- and last-layer outputs. - // Note: `last_block_out == final_out` — they alias the same tensor. - // Layer indices (mid_layer, last_layer) are needed to synthesize - // file names like "enc.block.21.out" because conf::named renames - // the tensor in place and the trailing "enc.final" rename - // overwrites the spot-check name. + // Mid- and last-layer spot-check outputs. last_block_out aliases + // final_out. Layer indices synthesize file names like + // "enc.block.21.out" because conf::named renames in place and the + // trailing "enc.final" rename overwrites the spot-check name. ggml_tensor * mid_block_out = nullptr; ggml_tensor * last_block_out = nullptr; int mid_block_idx = -1; @@ -124,10 +92,8 @@ struct EncoderBuild { ggml_tensor * mel_in = nullptr; // Sinusoidal positional embedding handle, ne=[d_model, 2*T_enc-1, 1, 1]. - // Sub-stages 3c+ need this; the driver computes the buffer - // host-side from T_enc (which is read from `out->ne[1]` after - // build) and uploads via ggml_backend_tensor_set. Null in 3a/3b - // because the FF1 sub-stage doesn't reference it. + // The driver computes the buffer host-side from T_enc (read from + // out->ne[1] after build) and uploads via ggml_backend_tensor_set. ggml_tensor * pos_emb_in = nullptr; // ChunkedLimited attention mask, ne=[T_enc, T_enc, 1, 1] f32. Null @@ -191,27 +157,20 @@ struct EncoderBuild { // mel input (the second dim of the row-major [n_mels, n_frames] // MelFrontend output) and shapes the input handle accordingly. // -// Returns an EncoderBuild with all three tensor handles populated. -// On any failure (insufficient mem_size, weights/hparams mismatch, -// degenerate frame count) the returned struct has nullptr fields and -// a diagnostic is logged via stderr; the caller must check. -// kv_type: GGML type for K/V activations in flash attention. +// Returns an EncoderBuild with all tensor handles populated. On failure +// (insufficient mem_size, weights/hparams mismatch, degenerate frame +// count) the struct has nullptr fields and a diagnostic is logged. +// kv_type: GGML type for K/V activations in flash attention; // GGML_TYPE_COUNT means "auto" (f16 for quantized weights, f32 for f32). -// backend_name: primary backend name (e.g. "MTL0", "Vulkan0", "CPU") for -// auto-detecting optimal conv strategy. Env vars override if set. -// Runtime override for chunked_limited_with_rc streaming. When the -// model declares enc_att_context_style=ChunkedLimitedWithRc and the -// caller wants to engage the chunked mask (i.e. buffered streaming on -// parakeet-unified-en-0.6b), pass a non-null pointer to a -// BufferedStreamMaskOverride with (L, C, R) in encoder frames. The -// builder then allocates a chunked_mask_in input tensor (the same -// shape the ChunkedLimited path uses) and routes the block params -// through the ChunkedLimited code path so rel_pos_mhsa picks up the -// precomputed mask. The caller fills the mask host-side via -// fill_chunked_limited_with_rc_mask after the compute buffer is -// allocated. Leaving this argument null keeps the offline behavior: -// the unified encoder runs with full attention (att_context_size -// [-1, -1] from the GGUF) and no mask. +// backend_name: primary backend name ("MTL0", "Vulkan0", "CPU") for +// auto-detecting conv strategy; env vars override. +// +// Runtime override for chunked_limited_with_rc streaming: pass a +// non-null BufferedStreamMaskOverride with (L, C, R) in encoder frames +// (buffered streaming on parakeet-unified-en-0.6b) to allocate a +// chunked_mask_in tensor and route blocks through the ChunkedLimited +// path; the caller fills the mask host-side. Null keeps offline full +// attention (att_context_size [-1, -1]). struct BufferedStreamMaskOverride { int left_frames; int chunk_frames; @@ -256,28 +215,24 @@ struct StreamingEncoderCacheIO { std::vector channel_out; std::vector time_out; - // Optional KV cache (driver-owned persistent tensors, one per - // layer, [d_model, T_cache] each): pre-projected attention - // keys/values replacing the last_channel recompute. When all four - // vectors are sized to n_layers and the builder enables the KV - // path (i.e. no dump harness active, since the numerical gate - // compares last_channel), blocks consume k_in/v_in, emit - // k_out/v_out rotation tensors, and leave channel_out null. + // Optional KV cache (driver-owned persistent tensors, one per layer, + // [d_model, T_cache] each): pre-projected attention keys/values + // replacing the last_channel recompute. When all four vectors are + // sized to n_layers and the builder enables the KV path, blocks + // consume k_in/v_in, emit k_out/v_out, and leave channel_out null. std::vector k_in; std::vector v_in; std::vector k_out; std::vector v_out; // Optional rel-pos projection memoization (driver-owned persistent - // tensors, one per layer, shape [head_dim, pos_len, n_head, 1]). - // pos_emb is a pure function of the chunk geometry, so - // attn_pos_w @ pos_emb is identical for every chunk with the same - // pos_len. When pos_proj_len matches the pos_len this build - // derives, the blocks consume pos_proj[i] directly and the graph - // references no pos_emb input at all (the builder leaves - // eb.pos_emb_in null). On mismatch the blocks compute the - // projection inline as before; the driver then refills the cache - // for the new geometry (see ensure_pos_proj_cache in model.cpp). + // tensors, one per layer, [head_dim, pos_len, n_head, 1]). pos_emb is + // a pure function of chunk geometry, so attn_pos_w @ pos_emb is + // identical for every chunk with the same pos_len. When pos_proj_len + // matches this build's pos_len, blocks consume pos_proj[i] directly + // and the graph references no pos_emb input (eb.pos_emb_in null); on + // mismatch blocks compute the projection inline and the driver + // refills the cache (see ensure_pos_proj_cache in model.cpp). std::vector pos_proj; int pos_proj_len = -1; }; diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index f2191e36..9c34cd5d 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -1,40 +1,13 @@ // arch/parakeet/model.cpp - Parakeet family handler. // -// Phase 4 step 1: load() now binds a runtime backend (Metal on Apple -// Silicon, CPU elsewhere or on Metal init failure). The second-stage -// gguf_init_from_file uses no_alloc=true so ggml builds the tensor -// catalog without touching the data section; we then allocate a -// backend buffer for every tensor via ggml_backend_alloc_ctx_tensors -// and stream the GGUF data section directly into it via -// ggml_backend_tensor_set. After load returns, every borrowed -// ggml_tensor* in `weights` has its data pointer rebound to backend -// memory. -// -// Stage 1 (header-only via transcribe::Loader, done by the dispatcher -// before we get here) is unchanged. Stage 2's gguf_context is freed -// before load() returns; the model's persistent state is ctx_meta + -// backend_buffer + backend_handle, freed in that order in the -// destructor. -// -// Phase 4 step 3a: run() now computes the mel front-end and runs -// the pre_encode subsampling stack of the encoder, with the result -// dumped via TRANSCRIBE_DUMP_DIR for the numerical accuracy harness. -// The remaining sub-stages (3b conformer FF1, 3c MHSA, 3d conv -// module, 3e FF2/norm_out, 3f loop over 24 blocks) progressively -// extend the encoder graph in arch/parakeet/encoder.cpp; the -// run-driver wiring here doesn't need to change. The decoder -// (predictor + joint + TDT) lands in phase 5. run() returns OK -// today, but every result accessor (full_text / segment / word / -// token) still returns its safe sentinel because nothing has -// populated the result yet. -// -// What's still missing relative to a fully-functional v1: -// - Conformer blocks (3b-3f). -// - Predictor + joint + TDT decode driver (phase 5). -// - Quantization. fp32 only today. -// - ggml_gallocr_t for repeat-call buffer reuse (phase 4 step 5). -// - Public timings API. t_load_us is captured on the model but -// not surfaced; phase 4 step 6 adds transcribe_timings. +// load() binds a runtime backend, opens the GGUF with no_alloc=true so +// ggml builds the tensor catalog without touching the data section, +// allocates a backend buffer for every tensor, and streams the data +// section into it; afterward every borrowed ggml_tensor* in `weights` +// points at backend memory. The persistent model state is ctx_meta + +// backend_buffer + backends, freed in that order in the destructor. +// run() computes the mel front-end, runs the encoder graph, then +// host-decodes (predictor + joint + TDT). #include "parakeet.h" @@ -56,12 +29,10 @@ #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" -// ggml-cpu.h no longer needed — threading is set via the registry -// proc address pattern, not ggml_backend_cpu_set_n_threads directly. #include "gguf.h" -// No backend-specific #includes needed — ggml's device registry in -// ggml-backend.h discovers Metal/Vulkan/CUDA/BLAS at runtime. +// No backend-specific #includes: ggml's device registry discovers +// Metal/Vulkan/CUDA/BLAS at runtime. #include #include @@ -82,22 +53,15 @@ namespace transcribe::parakeet { -// Forward declaration so the anonymous-namespace functions below can -// take the address of the registry entry. The full definition lives at -// the bottom of the file (after the per-call function definitions it -// references). +// Forward declaration; full definition at the bottom of the file. extern const Arch arch; -// Cheap insurance against a future contributor accidentally dropping the -// inheritance relationship. Cost: zero. Caught at compile time. static_assert(std::is_base_of_v); static_assert(std::is_base_of_v); ParakeetSession::~ParakeetSession() { - // Tear down per-call compute state. Order matters: the scheduler - // owns the compute buffers, the context owns the tensor metadata, - // both must be freed before the model's backend plan (which - // outlives the context, owned by ParakeetModel). + // Tear down per-call compute state before the model's backend plan + // (which outlives the context): scheduler, then context. if (sched != nullptr) { ggml_backend_sched_free(sched); sched = nullptr; @@ -108,12 +72,8 @@ ParakeetSession::~ParakeetSession() { } encoder_out = nullptr; - // Streaming cache tensors live in their own ggml_context with - // their own backend buffer (allocated lazily on first stream_begin - // for ChunkedLimited variants). Free in the same order as the - // weights buffer/session in ParakeetModel: backend buffer first (it - // may hold a backend ref), then the ggml_context that owned the - // tensor metadata. + // Streaming cache tensors live in their own ggml_context + backend + // buffer. Free buffer first (may hold a backend ref), then the ctx. if (stream_caches.buffer != nullptr) { ggml_backend_buffer_free(stream_caches.buffer); stream_caches.buffer = nullptr; @@ -141,15 +101,9 @@ ParakeetSession::~ParakeetSession() { } ParakeetModel::~ParakeetModel() { - // Teardown order: ctx_meta → backend_buffer → plan backends - // (reverse init order). - // - // ctx_meta owns the tensor metadata (ggml_tensor structs); the - // tensors' data pointers reference backend_buffer's interior. - // Freeing ctx_meta drops the metadata only — the buffer is - // independent. backend_buffer must be freed before the backends - // because the buffer was allocated against a backend and may - // hold a reference to it. Backends freed in reverse init order. + // Teardown order: ctx_meta → backend_buffer → plan backends. The + // buffer must be freed before the backends (it holds a backend ref); + // backends freed in reverse init order. if (bn_fused_ctx != nullptr) { ggml_free(bn_fused_ctx); bn_fused_ctx = nullptr; @@ -158,8 +112,8 @@ ParakeetModel::~ParakeetModel() { ggml_backend_buffer_free(bn_fused_buffer); bn_fused_buffer = nullptr; } - // The CPU conv_pw F32 promotion session + buffer (no-op on GPU primary - // backends; non-null only when promote_conv_pw_to_f32_on_cpu ran). + // CPU conv_pw F32 promotion ctx + buffer (non-null only when + // promote_conv_pw_to_f32_on_cpu ran). if (conv_pw_f32_ctx != nullptr) { ggml_free(conv_pw_f32_ctx); conv_pw_f32_ctx = nullptr; @@ -191,18 +145,10 @@ namespace { constexpr float kBnEps = 1e-5f; // Resolve the runtime language hint to a prompt-table index for the -// multilingual prompt-conditioned variants -// (NeMo's EncDecRNNTBPEModelWithPrompt). The dictionary carries both -// canonical BCP-47 codes ("en-US") and short aliases ("en") that map -// to the same index — we just do an exact-string lookup. -// -// Empty / null hint maps to the dictionary's `auto` slot (i.e. -// `prompt.auto_id`). When the model exposes language detection (its -// hparams declare auto_id), this is the "let the model emit a -// tag" path. Unknown hints return -1 (the caller should -// surface this as TRANSCRIBE_ERR_UNSUPPORTED_LANGUAGE — capability -// validation in transcribe.cpp will already have screened most -// invalid hints; this is the prompt-dictionary-vs-caps mismatch case). +// multilingual prompt-conditioned variants. The dictionary carries both +// BCP-47 codes ("en-US") and short aliases ("en"); exact-string lookup. +// Empty/null hint maps to the dictionary's auto slot (prompt.auto_id); +// unknown hints return -1 (caller surfaces UNSUPPORTED_LANGUAGE). int32_t resolve_prompt_id(const ParakeetHParams & hp, const char * language_hint) { @@ -210,9 +156,8 @@ int32_t resolve_prompt_id(const ParakeetHParams & hp, const bool empty_hint = (language_hint == nullptr) || (*language_hint == '\0'); if (empty_hint) { - // Auto-language slot. Models without an explicit auto entry - // leave auto_id at -1; in that case the dictionary's first - // entry is the conservative default. + // Auto-language slot; fall back to the dictionary's first entry + // when there is no explicit auto_id. if (hp.prompt_auto_id >= 0) return hp.prompt_auto_id; if (!hp.prompt_dictionary_indices.empty()) { return hp.prompt_dictionary_indices.front(); @@ -227,13 +172,10 @@ int32_t resolve_prompt_id(const ParakeetHParams & hp, return -1; } -// Fill a [num_prompts, T_enc, 1, n_batch] float buffer with a one-hot -// vector at column `prompt_id` of each utterance, replicated across -// the T_enc axis. The host-side replication keeps the in-graph -// step a single concat + two matmuls, with no in-graph one_hot / -// broadcast machinery. `prompt_ids` carries one id per utterance. -// -// Returns false if any per-utterance prompt_id is out of range. +// Fill a [num_prompts, T_enc, 1, n_batch] buffer with a one-hot at +// column prompt_id of each utterance, replicated across T_enc (one id +// per utterance in prompt_ids). Host-side replication keeps the in-graph +// step a single concat + two matmuls. Returns false if any id is OOR. bool fill_prompt_one_hot(std::vector & out, int num_prompts, int T_enc, @@ -262,15 +204,12 @@ bool fill_prompt_one_hot(std::vector & out, return true; } -// Fuse inference-time BatchNorm into precomputed scale + bias tensors. -// BN eval: y = (x - mean) / sqrt(var + eps) * weight + bias -// Fused: y = x * scale + shift -// where scale = weight / sqrt(var + eps) -// shift = bias - mean * scale -// -// Allocates fused tensors in a separate ggml context + CPU buffer on -// the model. Called once at load time; the fused tensors are then -// referenced by the encoder graph instead of the raw BN parameters. +// Fuse inference-time BatchNorm into scale + bias tensors: +// BN eval: y = (x - mean) / sqrt(var + eps) * weight + bias +// fused: y = x * scale + shift, scale = weight / sqrt(var + eps), +// shift = bias - mean * scale +// Fused tensors live in a separate ggml ctx + CPU buffer, computed once +// at load and referenced by the encoder graph in place of the raw BN. transcribe_status fuse_batch_norm(ParakeetModel & m) { const size_t n_blocks = m.weights.blocks.size(); if (n_blocks == 0) return TRANSCRIBE_OK; @@ -278,7 +217,6 @@ transcribe_status fuse_batch_norm(ParakeetModel & m) { const int64_t d = m.hparams.enc_d_model; const size_t tensor_bytes = static_cast(d) * sizeof(float); - // Create a context for the fused tensors (2 per block). const size_t ctx_size = n_blocks * 2 * ggml_tensor_overhead() + 256; ggml_init_params params = {ctx_size, nullptr, /*no_alloc=*/true}; m.bn_fused_ctx = ggml_init(params); @@ -291,9 +229,7 @@ transcribe_status fuse_batch_norm(ParakeetModel & m) { b.conv_bn_fused_bias = ggml_new_tensor_1d(m.bn_fused_ctx, GGML_TYPE_F32, d); } - // Allocate on the CPU backend. The plan's scheduler list always - // has CPU last as the fallback; for strict-CPU loads it is the - // only entry. + // Allocate on the CPU backend (always last in the scheduler list). m.bn_fused_buffer = ggml_backend_alloc_ctx_tensors( m.bn_fused_ctx, m.plan.scheduler_list.back()); if (m.bn_fused_buffer == nullptr) return TRANSCRIBE_ERR_BACKEND; @@ -322,26 +258,11 @@ transcribe_status fuse_batch_norm(ParakeetModel & m) { return TRANSCRIBE_OK; } -// On a CPU primary backend, dequantize the conformer 1×1 pointwise -// conv weights (pw1, pw2) from F16 back to F32. The shared quantizer -// (post commit 6fee9b9) routes Conformer ConvPw to F16 across all -// presets so GPU backends with native F16 compute (Metal, Vulkan, -// CUDA) get the bandwidth halving for free; on CPUs without native -// F16 arithmetic (Zen 2 and earlier) the per-element F16→F32 -// upconvert per matmul erases the bandwidth win and outweighs the -// original cost. -// -// Today's parakeet GGUFs predate the universal F16 conv_pw policy -// and ship F32 conv_pw weights, so this function is a no-op against -// the current on-disk artifacts. The wiring is in place so the next -// regen does the right thing automatically. See the same comment -// block on the cohere side for the cost trade and the ~235 MB of -// "wasted" originals that stay resident in the main weight buffer -// (ggml's backend buffer model does not support freeing individual -// tensors). -// -// The machinery itself lives in transcribe::load_common — this -// function is just the parakeet-specific weight walk. +// On a CPU primary backend, dequantize the conformer 1×1 pointwise conv +// weights (pw1, pw2) from F16 to F32: CPUs without native F16 arithmetic +// pay a per-matmul F16→F32 upconvert that erases the bandwidth win. The +// machinery lives in transcribe::load_common; this is the parakeet- +// specific weight walk. transcribe_status promote_conv_pw_to_f32_on_cpu(ParakeetModel & m) { std::vector slots; slots.reserve(m.weights.blocks.size() * 2); @@ -359,24 +280,14 @@ transcribe_status promote_conv_pw_to_f32_on_cpu(ParakeetModel & m) { } // Default variant string when the GGUF did not carry stt.variant. -// Defaulting belongs in the family handler, not the loader: each -// family knows its own canonical default and the loader does not. constexpr const char k_default_variant[] = "tdt-0.6b-v2"; // Allocate the streaming encoder caches (cache_last_channel, -// cache_last_time). Called lazily on the first -// stream_begin for ChunkedLimited variants. Layout matches NeMo's -// get_initial_cache_state exactly: see ParakeetStreamingCaches in -// parakeet.h for shapes and the per-variant sizing rule. -// -// The cache tensors live in their own ggml_context with a dedicated -// backend buffer, allocated on the model's primary backend so reads -// and writes inside the encoder graph stay backend-local. The buffer -// is freed in ~ParakeetSession. -// -// Initialization is idempotent — calling on an already-initialized -// caches struct is a no-op. To zero contents on a fresh stream, call -// zero_streaming_caches separately. +// cache_last_time), lazily on the first stream_begin for ChunkedLimited +// variants. Layout matches NeMo's get_initial_cache_state (see +// ParakeetStreamingCaches in parakeet.h). The tensors live in their own +// ggml_context + backend buffer on the primary backend (freed in the +// dtor). Idempotent; zero_streaming_caches clears contents separately. transcribe_status init_streaming_caches(ParakeetSession * pc, ParakeetModel * pm) { @@ -398,9 +309,7 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, return TRANSCRIBE_ERR_INVALID_ARG; } - // 4 tensors per layer (last_channel + last_time + last_k + last_v) - // plus headroom for descriptor / view tensors created during graph - // build. + // 4 tensors per layer (last_channel/time/k/v) plus graph-build headroom. const size_t ctx_size = static_cast(4 * n_layer + 8) * ggml_tensor_overhead(); ggml_init_params ip {}; @@ -470,11 +379,9 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, return TRANSCRIBE_OK; } -// Zero the streaming caches and reset cursors. Called at the start of -// every stream_begin so each utterance starts from a clean slate -// (NeMo's get_initial_cache_state returns zeros every time). -// -// The backend buffer survives — only the contents are cleared. +// Zero the streaming caches and reset cursors at each stream_begin +// (NeMo's get_initial_cache_state returns zeros every time). The backend +// buffer survives; only the contents are cleared. void zero_streaming_caches(ParakeetSession * pc) { if (pc->stream_caches.buffer != nullptr) { ggml_backend_buffer_clear(pc->stream_caches.buffer, 0); @@ -484,10 +391,8 @@ void zero_streaming_caches(ParakeetSession * pc) { pc->stream_caches.pcm_start_sample = 0; } -// Reset the streaming decoder state (LSTM h/c, prev token, frame -// offset). Sized to the model's predictor on first call; resized -// in-place on subsequent calls if the predictor reshape ever happens -// (it won't, but the guard is cheap). +// Reset the streaming decoder state (LSTM h/c, prev token, frame offset) +// to the model's predictor sizing. void reset_streaming_decoder_state(ParakeetSession * pc, const ParakeetModel * pm) { @@ -514,59 +419,36 @@ transcribe_status load( const transcribe_model_load_params * params, transcribe_model ** out_model) { - // The dispatcher has already verified out_model is non-null and - // the loader has a valid gguf_context with general.architecture - // set. *out_model is currently null (the dispatcher cleared it). - // Backend and device selection are resolved below via load_common. + // The dispatcher has verified out_model is non-null and cleared it, + // and the loader has a valid gguf_context. const int64_t t_load_start = ggml_time_us(); auto m = std::make_unique(); m->arch = &arch; - m->t_load_us = 0; // will be set just before the successful return + m->t_load_us = 0; - // Variant defaulting belongs here, not in the loader: each family - // knows its own canonical default. We accept the GGUF's stt.variant - // verbatim if present and only fall back when it is empty. - // - // The variant string is descriptive metadata for users — it - // surfaces through transcribe_model_variant_string but does NOT - // drive any behavior decision. Capability differences between - // Parakeet variants (v2 English vs v3 multilingual, etc.) are - // expressed as stt.capability.* and general.languages KV, read - // below. The only PLAN.md-blessed exception is the decoder-kind - // prefix dispatch for hybrid models, which Parakeet v1 does not - // exercise (only TDT ships). + // Variant defaulting belongs in the family handler. The variant + // string is descriptive metadata (surfaced via + // transcribe_model_variant_string) and drives no behavior decision; + // capability differences are expressed as stt.capability.* / + // general.languages KV, read below. if (loader.variant().empty()) { m->variant = k_default_variant; } else { m->variant = loader.variant(); } - // backend is set below, after the runtime backend is initialized. - // Until then m->backend stays empty per the public ABI semantic - // for "no backend currently bound". + // Empty until the runtime backend is bound below (public ABI: "no + // backend currently bound"). m->backend.clear(); - // Capability resolution, KV-driven. Order matters: - // - // 1. apply_family_invariants populates the family defaults - // (native_sample_rate, supports_translate). These are the - // "the architecture supplies a default" half of PLAN.md's - // capability precedence rule. - // 2. read_capability_kv overlays any stt.capability.* KV the - // converter wrote. This is the "GGUF KV is authoritative" - // half. Absent KV leaves the family default in place; - // present-but-wrong-type KV is a converter bug and propagates - // as TRANSCRIBE_ERR_GGUF. - // - // Information-gap default for the languages chain. read_languages_kv - // below leaves these in place if general.languages is absent — the - // result is the public ABI's "n_languages == 0, languages == nullptr" - // observation, documented in include/transcribe.h as "we don't - // know" rather than "the model has no languages". A v3 multilingual - // GGUF supplies the real list via general.languages and - // read_languages_kv overwrites the defaults via set_languages. + // Capability resolution, KV-driven. apply_family_invariants + // populates the family defaults; read_capability_kv then overlays any + // stt.capability.* KV (absent leaves the default; wrong-type + // propagates as TRANSCRIBE_ERR_GGUF). n_languages=0 / languages=null + // is the "we don't know" state until read_languages_kv overwrites it + // from general.languages. apply_family_invariants(*m); m->caps.n_languages = 0; m->caps.languages = nullptr; @@ -577,48 +459,35 @@ transcribe_status load( return st; } - // Languages from general.languages, installed via set_languages - // (which keeps the language pointer chain in caps in sync). The - // strings are copied into the model so the loader's gguf_context - // can be freed after this point without invalidating them. + // Languages from general.languages (copied into the model, so the + // loader's gguf_context can be freed afterward). if (const transcribe_status st = read_languages_kv(loader.gguf(), *m); st != TRANSCRIBE_OK) { return st; } - // Tokenizer ingest. Reads tokenizer.ggml.* keys from the loader's - // gguf_context and copies the strings / scores / token types into - // the Tokenizer's own storage. Like read_languages_kv above, this - // does not retain a borrow into the gguf_context. + // Tokenizer ingest: copies tokenizer.ggml.* into the Tokenizer's own + // storage (no borrow into the gguf_context). if (const transcribe_status st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) { return st; } // Architecture KV (encoder / predictor / joint / frontend dims). - // Read directly from the loader's gguf_context — these are - // hparams, not tensor data, so we can do this before the second - // open. read_parakeet_hparams enforces cross-field invariants - // (d_model divisible by n_heads, etc.) so we know the shapes - // we're about to validate against are themselves consistent. + // read_parakeet_hparams enforces cross-field invariants so the shapes + // we validate against are consistent. if (const transcribe_status st = read_parakeet_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) { return st; } - // Derive supports_streaming from hparams. - // - // ChunkedLimited + (L, R) >= 0 — cache-aware streaming - // (nemotron-speech-streaming-en-0.6b). The cache_last_channel / - // cache_last_time path is engaged at stream_begin. + // Derive supports_streaming from hparams: + // ChunkedLimited + (L, R) >= 0 — cache-aware streaming + // (nemotron-speech-streaming-en-0.6b). // ChunkedLimitedWithRc + non-empty (L, C, R) menu — buffered - // streaming (parakeet-unified-en-0.6b). The buffered driver - // re-runs the offline encoder per chunk with a 3-tuple chunk - // mask; no per-layer cache. - // - // Offline parakeet variants (Regular full attention or no chunking) - // stay non-streaming regardless of any KV claim. + // streaming (parakeet-unified-en-0.6b). + // Offline variants stay non-streaming. const bool cache_aware_streaming = (m->hparams.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited) && @@ -630,20 +499,14 @@ transcribe_status load( !m->hparams.enc_att_chunk_left_choices.empty() && !m->hparams.enc_att_chunk_chunk_choices.empty() && !m->hparams.enc_att_chunk_right_choices.empty(); - // supports_streaming is the generic gate; the configurable - // streaming geometry (the (left, chunk, right) menu for buffered, - // the att_context_right menu for cache-aware) is reached via the - // parakeet stream extensions, not advertised as flat caps fields. + // supports_streaming is the generic gate; the configurable geometry + // is reached via the parakeet stream extensions, not flat caps. if (buffered_streaming || cache_aware_streaming) { m->caps.supports_streaming = true; } - // Construct the mel front-end now that hparams are available. - // The MelFrontend constructor precomputes the periodic Hann - // window + Slaney mel filterbank (~50 ms on jfk-sized configs); - // doing it here amortizes the cost across every transcribe_run - // on every context derived from this model. The instance is - // const-after-construction and thread-safe per its contract. + // Construct the mel front-end (precomputes Hann window + Slaney mel + // filterbank); amortized across every run on every context. { transcribe::MelConfig cfg {}; cfg.sample_rate = m->hparams.fe_sample_rate; @@ -655,31 +518,18 @@ transcribe_status load( cfg.f_min = m->hparams.fe_f_min; cfg.f_max = m->hparams.fe_f_max; cfg.normalize = m->hparams.fe_normalize; - // NeMo's FilterbankFeatures.stft uses pad_mode="constant" for - // the STFT center-pad (see features.py line 385). The cpp - // MelConfig default is "reflect" — matching cpp's behavior to - // NeMo here. The two differ in the first/last few STFT frames: - // reflect pads with mirrored audio, constant pads with zeros. - // Offline this is one boundary per utterance and the residual - // washes it out; streaming sees boundary effects every chunk - // and the per-feature normalize amplifies the divergence into - // the encoder. + // NeMo's FilterbankFeatures.stft uses pad_mode="constant" (the + // cpp MelConfig default is "reflect"). The two differ in the + // first/last few STFT frames; offline the residual washes out but + // streaming sees it every chunk, amplified by per-feature normalize. cfg.pad_mode = "constant"; m->mel.emplace(cfg); } - // Stage 2: reopen the file with no_alloc=true + session=&ctx_meta. - // ggml builds the tensor catalog (metadata only — no data - // buffers); we then bind a runtime backend, allocate a backend - // buffer for every tensor in ctx_meta, and stream the GGUF data - // section directly into that backend buffer. After this - // sequence, ggml_get_tensor(ctx_meta, name)->data points at - // backend memory, not host malloc. - // - // The first gguf_context held by the loader stays alive until - // load() returns (the dispatcher owns its lifetime); we don't - // touch it here. The second gguf_context (gguf_data below) is - // freed before load() returns. + // Reopen the file with no_alloc=true + ctx=&ctx_meta: ggml builds the + // tensor catalog (metadata only), then we bind a backend, allocate a + // buffer for every tensor, and stream the data section into it. The + // second gguf_context (gguf_data) is freed before load() returns. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -687,16 +537,13 @@ transcribe_status load( gguf_context * gguf_data = gguf_init_from_file(loader.path().c_str(), init_params); if (gguf_data == nullptr) { - // ggml has already logged a structured error. The first-stage - // open succeeded so this is genuinely unexpected — most - // likely a race where the file changed between the two opens, - // or a permissions transition. Surface as ERR_GGUF. + // The first-stage open succeeded, so this is unexpected (file + // changed between opens, or a permissions transition). return TRANSCRIBE_ERR_GGUF; } - // Validate every tensor against the canonical catalog. find_tensor - // only inspects type + shape (not data pointers), so this works - // before the backend buffer is bound. + // Validate every tensor against the canonical catalog (type + shape + // only, so it works before the backend buffer is bound). if (const transcribe_status st = build_parakeet_weights(m->ctx_meta, m->hparams, m->weights); st != TRANSCRIBE_OK) @@ -705,13 +552,8 @@ transcribe_status load( return st; } - // Resolve the backend plan via ggml's device registry. The - // caller's requested backend (AUTO, CPU, METAL, VULKAN) is - // honored here; see transcribe-backend.h + transcribe-load-common.h - // for the full semantic. This is runtime-dynamic: the library - // code has no compile-time #if TRANSCRIBE_METAL / #if - // TRANSCRIBE_VULKAN guards. Whatever backends ggml was compiled - // with are discovered here. + // Resolve the backend plan via ggml's device registry (runtime, no + // compile-time backend guards). See transcribe-load-common.h. const transcribe_backend_request backend_req = (params != nullptr) ? params->backend : TRANSCRIBE_BACKEND_AUTO; @@ -723,15 +565,11 @@ transcribe_status load( return st; } - // Label for the public API: report the primary backend. m->backend = ggml_backend_name(m->plan.primary); m->primary_backend = m->plan.primary; // Allocate a backend buffer for every tensor in ctx_meta on the - // primary backend. After this returns, each ggml_tensor in - // ctx_meta has its `buffer` and `data` slot bound to the backend - // allocation. We still need to upload the actual weight bytes - // from the GGUF data section. + // primary backend; the weight bytes are streamed in below. ggml_backend_buffer_t weights_buffer = ggml_backend_alloc_ctx_tensors(m->ctx_meta, m->plan.primary); if (weights_buffer == nullptr) { @@ -744,11 +582,9 @@ transcribe_status load( ggml_backend_buffer_set_usage(weights_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - // Stream tensor data from the GGUF file into the backend buffer - // slots. See transcribe-load-common.h for the shared loop; it - // uses ggml_backend_tensor_set so this works for both host- - // memory backends (CPU, Metal on Apple Silicon) and discrete - // GPUs. + // Stream tensor data from the GGUF into the backend buffer slots + // (shared loop in transcribe-load-common.h; works on host-memory + // backends and discrete GPUs). if (const transcribe_status st = transcribe::load_common::stream_tensor_data( loader.path(), gguf_data, m->ctx_meta, "parakeet"); st != TRANSCRIBE_OK) @@ -757,17 +593,11 @@ transcribe_status load( return st; } - // gguf_data only carried tensor offsets + names; both have been - // consumed by build_parakeet_weights and the streaming loop - // above. m->ctx_meta + m->backend_buffer + m->backend_handle + - // m->weights are now the model's entire state. gguf_free(gguf_data); - // Fuse BatchNorm parameters into scale + bias for the encoder - // graph. This replaces 4 elementwise ops per block with 2. Skipped - // for streaming variants whose conv module uses LayerNorm (raw - // bn.weight / bn.bias travel through the graph as LN scale/bias, - // and there are no running stats to fuse against). + // Fuse BatchNorm into scale + bias (replaces 4 elementwise ops per + // block with 2). Skipped for LayerNorm conv-module variants (no + // running stats to fuse against). if (m->hparams.enc_conv_norm_type == ParakeetHParams::ConvNormType::BatchNorm) { @@ -778,37 +608,25 @@ transcribe_status load( } } - // On CPU primary backend, dequantize conv pointwise weights back - // to F32 — Zen 2 class CPUs don't have native F16 compute and the - // upconvert cost per matmul outweighs the bandwidth savings from - // the smaller weight. No-op on GPU backends and on parakeet GGUFs - // that predate the universal F16 conv_pw quantizer policy (today, - // all of them). + // On CPU primary backend, dequantize conv pointwise weights to F32 + // (see promote_conv_pw_to_f32_on_cpu). No-op on GPU backends. if (const transcribe_status st = promote_conv_pw_to_f32_on_cpu(*m); st != TRANSCRIBE_OK) { return st; } - // Phase 5: build the host mirror of the predictor + joint - // weights. This is a one-shot copy from backend memory into - // std::vectors on the model. The decoder runs on host - // (see decoder.h for the rationale), and centralizing the - // backend → host copy here means the per-call decode loop - // pays no readback cost. + // Build the host mirror of the predictor + joint weights (one-shot + // backend → host copy; the decoder runs on host, see decoder.h). if (const transcribe_status st = build_host_decoder_weights(*m, m->host_decoder); st != TRANSCRIBE_OK) { return st; } - // Capture wall-clock load time on the model. The accumulator - // is on the base transcribe_model so the public timings API - // can read it without per-family casts. m->t_load_us = ggml_time_us() - t_load_start; - // Hand off to the caller. From here on the caller owns the model - // and must call transcribe_model_free. + // The caller now owns the model and must call transcribe_model_free. *out_model = m.release(); return TRANSCRIBE_OK; } @@ -818,12 +636,8 @@ transcribe_status init_context( const transcribe_session_params * params, transcribe_session ** out_ctx) { - // The central dispatcher has already null-checked model, params, - // and out_ctx. We still want to be sure the model we received is - // actually a ParakeetModel — `arch` is the discriminator. In debug - // builds an assert would be louder, but a routing mistake should - // not silently corrupt anything in release either, so we - // defensively check. + // The dispatcher has null-checked model/params/out_ctx; verify the + // model is actually a ParakeetModel (arch is the discriminator). if (model->arch != &arch) { return TRANSCRIBE_ERR_INVALID_ARG; } @@ -837,38 +651,12 @@ transcribe_status init_context( return TRANSCRIBE_OK; } -// Internal inference helper. Assumes pc/pm are non-null and the -// scheduler plan is populated; callers (run() and stream_finalize) -// validate that up front. n_samples must be > 0; pcm must be non-null. -// -// Does NOT call pc->clear_result(); callers handle result-snapshot -// management. (The public run() entry clears once before reaching here; -// stream_finalize relies on the dispatcher's clear at stream_begin time -// and does not re-clear, so its accumulated stream cursors survive.) -// -// Phase 4a — run/stream parity: -// -// For cache-aware-streaming Parakeet variants (today: nemotron- -// speech-streaming-en-0.6b), transcribe_run and the streaming hooks -// are siblings of this helper; neither owns the inference. That -// gives the streaming-of-whole path guaranteed-by-construction -// parity with one-shot run on the same audio. When real per-chunk -// encoder feeding lands (Phase 4b: cache_last_channel / -// cache_last_time / RNNT predictor carry), the streaming hooks -// will diverge from this helper and an empirical stream-vs-batch -// parity test becomes meaningful. -// Host-side TDT/RNNT/CTC decode + public result-hierarchy build for a -// single utterance's encoder output. Shared by the single-shot path -// True for a multilingual language-tag SentencePiece piece of the form -// "" (lowercase 2-3 letter language, '-', 2-4 letter region), e.g. -// "", "", "". The nemotron-3.5 multilingual vocab emits -// one of these per segment to mark the (detected or requested) language. We -// drop them from the public token/word/segment/text result by default so the -// transcript stays clean and word/timestamp aggregation isn't polluted by a -// zero-width tag "word". Requiring the '-REGION' part avoids matching control -// pieces like "". No-op for the English parakeets (no such pieces in -// their 1024-token vocab). Preserved when keep_special_tags is set -// (CLI --raw-tokens). +// True for a multilingual language-tag SentencePiece piece "" +// (2-3 letter language, '-', 2-4 letter region), e.g. "". The +// nemotron-3.5 vocab emits one per segment to mark the language; we drop +// them from the public result by default (clean transcript, uncorrupted +// word/timestamp aggregation). Requiring '-REGION' avoids matching +// control pieces like "". No-op for the English parakeets. static bool is_lang_tag_piece(const std::string & p) { const size_t n = p.size(); if (n < 7 || p.front() != '<' || p.back() != '>') return false; // min "" @@ -887,22 +675,15 @@ static bool is_lang_tag_piece(const std::string & p) { return i == end; // interior consumed exactly up to '>' } -// Single source of truth for "drop this piece from the public result when -// keep_special_tags is off." A piece is stripped if the vocab marks it -// CONTROL (canary/nemotron tag pieces are CONTROL-typed in the GGUF) or it -// matches the multilingual locale-tag pattern (transitional fallback -// for GGUFs predating the converter's CONTROL marking). Used by both the -// offline (decode_and_populate) and streaming (rebuild_streaming_result_text) -// result builders so they strip exactly the same set. +// Drop this piece from the public result when keep_special_tags is off: +// stripped if CONTROL-typed or matching the locale-tag pattern +// (transitional fallback). Shared by the offline and streaming builders. static bool is_strippable_special(const transcribe::Tokenizer & tok, int id) { return tok.is_control(id) || is_lang_tag_piece(tok.token(id)); } -// Normalize a decoded transcript's whitespace: collapse runs of ASCII spaces -// to one and trim both ends. Stripping an interior tag leaves its neighbours' -// spaces adjacent (double space); a stripped trailing tag leaves a trailing -// space. An ASR transcript never carries meaningful double/edge whitespace, so -// this is a no-op when nothing was stripped. Shared by both result builders. +// Collapse runs of ASCII spaces to one and trim both ends — cleans up the +// double/edge spaces a stripped tag leaves behind. Shared by both builders. static void normalize_transcript_whitespace(std::string & s) { std::string norm; norm.reserve(s.size()); @@ -918,15 +699,10 @@ static void normalize_transcript_whitespace(std::string & s) { s.swap(norm); } -// Milliseconds spanned by one encoder frame, following NeMo's reference -// conversion (collections/asr/parts/utils/timestamp_utils.py: -// time_seconds = frame_offset * window_stride * subsampling_factor, where -// window_stride is the preprocessor hop in seconds = hop_length / -// sample_rate). NeMo keeps the product in floating point; we keep the per- -// frame stride as a double and round once at the call site to the public -// API's ms grain. Returned as a double (NOT pre-rounded to an integer) so a -// non-integral stride doesn't quantize before the multiply. Single source of -// truth for both the offline and streaming result builders. +// Milliseconds per encoder frame (NeMo: subsampling_factor * hop_length / +// sample_rate). Returned as a double — not pre-rounded — so a non-integral +// stride doesn't quantize before the call-site multiply. Shared by both +// result builders. static double parakeet_ms_per_enc_frame(const ParakeetHParams & hp) { return 1000.0 * static_cast(hp.enc_subsampling_factor) * @@ -934,10 +710,10 @@ static double parakeet_ms_per_enc_frame(const ParakeetHParams & hp) { static_cast(hp.fe_sample_rate); } -// (run_one_shot_inner) and the batched path (run_batch_inner). `enc` is -// the row-major [T_enc, d_enc] encoder activation for ONE utterance; -// utt_index >= 0 tags the optional tensor dump per utterance, -1 for the -// single-shot name. Writes the session scratch result slot (tokens / +// Host-side TDT/RNNT/CTC decode + public result-hierarchy build for one +// utterance's encoder output. `enc` is the row-major [T_enc, d_enc] +// activation for ONE utterance; utt_index >= 0 tags the per-utterance +// dump (-1 for single-shot). Writes the session result slot (tokens / // words / segments / full_text / result_kind / has_result). static transcribe_status decode_and_populate( ParakeetSession * pc, @@ -959,9 +735,7 @@ static transcribe_status decode_and_populate( if (utt_index >= 0) { dump_name += ".b" + std::to_string(utt_index); } - // Optional dump of the encoder output as the decoder sees it, - // so the bring-up loop can verify the readback is faithful - // before chasing decoder bugs. + // Optional dump of the encoder output as the decoder sees it. if (transcribe::debug::enabled()) { const long long shape[2] = { T_enc, d_enc }; const char * stage = (enc_dump_name_override != nullptr && @@ -996,39 +770,27 @@ static transcribe_status decode_and_populate( } pc->t_decode_us = ggml_time_us() - t_dec_start; - // ----- Build the public result hierarchy ---------------------- + // ----- Build the public result hierarchy ----- // - // Conversion: each emitted token's `step_at_emit` is an encoder - // frame index. The encoder downsamples by `subsampling_factor` - // relative to the mel hop, so one encoder frame corresponds to - // `subsampling_factor * hop_length / sample_rate` seconds of - // audio. For Parakeet 0.6B v2/v3 that's 8 * 160 / 16000 = 0.08 s - // per frame, i.e. 80 ms. Shared with the streaming builder. + // step_at_emit is an encoder frame index; one frame is + // subsampling_factor * hop_length / sample_rate seconds (80 ms on + // v2/v3). Shared with the streaming builder. const double frame_to_ms = parakeet_ms_per_enc_frame(pm->hparams); - // SentencePiece word-boundary marker U+2581 ("▁"). UTF-8 bytes - // 0xE2 0x96 0x81. A token whose decoded text starts with the - // marker opens a new word; otherwise it extends the current - // word. We use the raw NeMo piece (not the post-decode "▁ → - // space" form) because the post-decode form leads with a space - // and that's a fragile thing to compare against — the raw - // piece's leading 3 UTF-8 bytes are unambiguous. + // SentencePiece word-boundary marker U+2581 ("▁", UTF-8 E2 96 81): a + // raw piece starting with it opens a new word. We use the raw NeMo + // piece (its leading 3 bytes are unambiguous), not the post-decode + // "▁ → space" form. constexpr const char k_sp_marker[] = "\xE2\x96\x81"; constexpr int k_sp_marker_len = 3; const transcribe::Tokenizer & tok = pm->tok; pc->tokens.reserve(pc->raw_tokens.size()); - // Strip multilingual language tags from the public result by - // default (clean transcript + uncorrupted word/timestamp aggregation). - // Gated on the standard keep_special_tags run param (CLI --raw-tokens) — - // the same switch canary/sensevoice use for their <|...|> tags. - // - // Primary detection is Tokenizer::is_control: convert-parakeet.py marks the - // tag pieces CONTROL in the GGUF token_type (read from the vocab shape, not - // a hard-coded list). The is_lang_tag_piece() pattern is a transitional - // fallback for GGUFs produced before that converter change; it can be - // dropped once all shipped artifacts carry the metadata. + // Strip multilingual language tags by default (gated on + // keep_special_tags / CLI --raw-tokens). Detection is + // Tokenizer::is_control (CONTROL token_type); is_lang_tag_piece is a + // transitional fallback for GGUFs predating that converter change. const bool strip_tags = (params == nullptr) ? true : !params->keep_special_tags; for (const TdtToken & rt : pc->raw_tokens) { if (strip_tags && is_strippable_special(tok, rt.id)) { @@ -1039,42 +801,28 @@ static transcribe_status decode_and_populate( te.p = rt.p; te.t0_ms = static_cast( std::llround(frame_to_ms * static_cast(rt.step_at_emit))); - // Per-token duration: clamp the duration_frames=0 case to a - // visually-distinct minimum of 0 ms (== t0). The result is a - // "point in time" token with no width. + // duration_frames=0 yields a zero-width "point in time" token. const double end_frame = static_cast(rt.step_at_emit) + static_cast(rt.duration_frames); te.t1_ms = static_cast(std::llround(frame_to_ms * end_frame)); - // Decode just this single token to get its visible text - // fragment. The decoder strips the SentencePiece marker via - // Tokenizer::decode (▁ → space), so the leading character - // for word-starting tokens is an ASCII space. te.text = tok.decode(&rt.id, 1); - // seg_index / word_index filled below. pc->tokens.push_back(std::move(te)); } - // Word + segment construction. v1 produces a single segment - // covering the entire clip; words are split on SentencePiece - // marker boundaries. Empty result → no segment, no words, no - // text — exactly what the public accessors return as their - // safe-sentinel state when there are no tokens. + // Word + segment construction. v1 produces a single segment over the + // whole clip; words split on SentencePiece marker boundaries. Empty + // result → no segment/words/text (the accessors' safe-sentinel state). if (!pc->tokens.empty()) { - // Single segment. transcribe_session::SegmentEntry seg; seg.t0_ms = pc->tokens.front().t0_ms; seg.t1_ms = pc->tokens.back().t1_ms; seg.first_token = 0; seg.n_tokens = static_cast(pc->tokens.size()); seg.first_word = 0; - // n_words filled after we count words below. - // Walk the raw token ids again to detect word boundaries. - // The first token always opens word 0; a token whose raw - // piece begins with the SentencePiece marker opens a new - // word. Continuation tokens (no marker) extend the current - // word. + // First token opens word 0; a raw piece beginning with the + // marker opens a new word, others extend the current word. transcribe_session::WordEntry cur_word; bool cur_word_open = false; @@ -1101,8 +849,7 @@ static transcribe_status decode_and_populate( if (starts_word) { open_new_word(static_cast(i), tk); } - // Whether we just opened a new word or not, the current - // word now contains this token. seg/word back-pointers: + // seg/word back-pointers. pc->tokens[i].seg_index = 0; pc->tokens[i].word_index = static_cast(pc->words.size()); } @@ -1113,11 +860,8 @@ static transcribe_status decode_and_populate( pc->words.push_back(std::move(cur_word)); } - // Materialize each word's text via the tokenizer (so the - // SentencePiece "▁ → space" substitution runs on the whole - // span at once, including any mid-word continuation - // pieces). The leading space from a word-opener token is - // trimmed: a "word" should not include its own leading + // Materialize each word's text via the tokenizer (so "▁ → space" + // runs over the whole span), trimming the word-opener's leading // space. std::vector id_buf; for (auto & wd : pc->words) { @@ -1148,32 +892,16 @@ static transcribe_status decode_and_populate( pc->full_text = std::move(full); pc->segments.push_back(std::move(seg)); - // Parakeet TDT produces token-level timestamps from the - // encoder frame indices (each emitted token's - // step_at_emit). Word and segment timestamps are derived - // by aggregating contained tokens. TOKEN is therefore the - // family's max_timestamp_kind, and everything above has - // already been populated at TOKEN granularity. - // - // Clamp to the caller's requested ceiling. AUTO resolves to - // the family max (TOKEN). A coarser request elides the - // finer levels: WORD drops the token list, SEGMENT drops - // tokens+words, NONE drops tokens+words and zeros the - // segment's t0/t1 so nothing dresses up as alignment data - // the caller did not ask for. - // - // The dispatcher has already rejected any request finer - // than TOKEN, so after the AUTO resolution below, eff is in - // {NONE, SEGMENT, WORD, TOKEN}. + // Clamp to the caller's requested ceiling. AUTO resolves to the + // family max (TOKEN); a coarser request elides the finer levels. + // The dispatcher has already rejected any request finer than + // TOKEN, so eff ends up in {NONE, SEGMENT, WORD, TOKEN}. transcribe_timestamp_kind eff = params->timestamps; if (eff == TRANSCRIBE_TIMESTAMPS_AUTO) { eff = pm->caps.max_timestamp_kind; // = TOKEN for parakeet } if (eff == TRANSCRIBE_TIMESTAMPS_NONE) { - // Keep the segment and its text, but drop alignment - // data. Tokens + words are a finer granularity than - // the caller asked for; segment timings are alignment - // too, so zero them. + // Keep the segment + text, drop all alignment data. pc->tokens.clear(); pc->words.clear(); auto & s = pc->segments.back(); @@ -1193,13 +921,9 @@ static transcribe_status decode_and_populate( s.first_token = 0; s.n_tokens = 0; } else if (eff == TRANSCRIBE_TIMESTAMPS_WORD) { - // Keep segment + word timings, drop the token table. - // Every back-reference into the cleared token table - // must also be zeroed so a caller iterating words or - // the segment can never index into a now-empty - // pc->tokens. The word_* / segment_* accessors return - // safe sentinels when n_tokens == 0, which is the - // documented behavior for "coarser than requested." + // Keep segment + word timings, drop the token table. Zero the + // back-references so word/segment accessors return safe + // sentinels (n_tokens == 0) rather than indexing it. pc->tokens.clear(); for (auto & w : pc->words) { w.first_token = 0; @@ -1225,23 +949,14 @@ static transcribe_status run_one_shot_inner( int n_samples, const transcribe_run_params * params) { - // Pre-run abort check. Parakeet is non-chunked today; this is the - // single observation point. A caller that wants to veto a run - // without paying encoder cost flips the callback's state and the - // next transcribe_run short-circuits here. + // Pre-run abort check (the single observation point on this path). if (pc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } - // Initialize the debug dumper from TRANSCRIBE_DUMP_DIR. Idempotent - // — only the first call has effect. transcribe::debug::init(); - // ----- Mel front-end ------------------------------------------- - // - // The MelFrontend was constructed once at load() time and lives - // on the model. compute() is documented thread-safe across - // contexts since the instance is const-after-construction. + // ----- Mel front-end ----- if (!pm->mel.has_value()) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet run: model has no MelFrontend (load skipped?)"); @@ -1262,29 +977,19 @@ static transcribe_status run_one_shot_inner( } pc->t_mel_us = ggml_time_us() - t_mel_start; - // ----- Reset per-call compute state ---------------------------- - // - // Free the previous run's compute_ctx (tensor metadata + cgraph). - // The gallocr persists across calls — it reuses its internal - // buffer when the graph topology is unchanged (same audio length) - // and reallocates automatically on topology change. The - // encoder_out borrowed pointer is invalidated here. + // ----- Reset per-call compute state ----- + // Free the previous run's compute_ctx; encoder_out is invalidated. if (pc->compute_ctx != nullptr) { ggml_free(pc->compute_ctx); pc->compute_ctx = nullptr; } pc->encoder_out = nullptr; - // ----- Build the compute context + encoder graph --------------- + // ----- Build the compute context + encoder graph ----- // - // The compute context only holds tensor metadata + the cgraph; - // tensor data is allocated separately by - // ggml_backend_alloc_ctx_tensors below. The full 24-block - // encoder graph has ~2200 tensors plus an 8192-slot cgraph - // (see encoder.cpp's ggml_new_graph_custom call). Empirically - // this uses ~1.23 MB of metadata arena (logged once during - // step 5 polish). 4 MB gives ~3x headroom for future per-block - // op additions (BN fold, gallocr metadata, etc.). + // compute_ctx holds tensor metadata + the cgraph (data is allocated + // by the scheduler below). The 24-block encoder graph uses ~1.23 MB + // of metadata arena; 4 MB gives ~3x headroom. { ggml_init_params init_params {}; init_params.mem_size = 4 * 1024 * 1024; @@ -1298,9 +1003,8 @@ static transcribe_status run_one_shot_inner( } } - // Resolve kv_type from the public enum to ggml_type. - // GGML_TYPE_COUNT is the sentinel for "auto" inside the encoder. - ggml_type resolved_kv = GGML_TYPE_COUNT; // auto + // GGML_TYPE_COUNT is the encoder's "auto" sentinel. + ggml_type resolved_kv = GGML_TYPE_COUNT; if (pc->kv_type == TRANSCRIBE_KV_TYPE_F32) resolved_kv = GGML_TYPE_F32; if (pc->kv_type == TRANSCRIBE_KV_TYPE_F16) resolved_kv = GGML_TYPE_F16; @@ -1308,16 +1012,11 @@ static transcribe_status run_one_shot_inner( pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, pm->backend.c_str()); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { - // build_encoder_graph already logged the diagnostic. - return TRANSCRIBE_ERR_GGUF; + return TRANSCRIBE_ERR_GGUF; // build_encoder_graph logged } - // ----- Allocate compute tensors via scheduler -------------------- - // - // The multi-backend scheduler dispatches ops to the best backend - // (GPU for matmuls if available, BLAS for CPU matmuls, CPU for - // the rest). It also manages compute buffer allocation with - // live-range packing. Created lazily, persists across calls. + // ----- Allocate compute tensors via the multi-backend scheduler. + // Created lazily, persists across calls. ----- if (pc->sched == nullptr) { pc->sched = ggml_backend_sched_new( pm->plan.scheduler_list.data(), nullptr, @@ -1336,25 +1035,17 @@ static transcribe_status run_one_shot_inner( return TRANSCRIBE_ERR_GGUF; } - // Upload the mel into the input tensor. The C++ MelFrontend - // produces a row-major [num_mels, n_frames] buffer, which is - // byte-identical to the ggml ne=[n_frames, num_mels, 1, 1] - // layout the encoder builder created (see encoder.cpp's layout - // cheat sheet for the derivation). + // Upload the mel; the row-major [num_mels, n_frames] buffer is + // byte-identical to the ggml ne=[n_frames, num_mels, 1, 1] input. ggml_backend_tensor_set(eb.mel_in, pc->mel_buf.data(), 0, pc->mel_buf.size() * sizeof(float)); - // Optional dump of the mel input for the numerical accuracy - // harness. Gated on TRANSCRIBE_DUMP_DIR; zero-cost when unset. transcribe::debug::dump_tensor( "enc.mel.in", eb.mel_in, "encoder.mel"); - // Prompt one-hot input (multilingual variants only). The graph - // builder allocates `prompt_one_hot_in` of shape - // [num_prompts, T_enc, 1, 1] when hp.has_prompt is true; the - // host fills it with the resolved language's one-hot replicated - // across T_enc frames. See resolve_prompt_id for the language - // string → index lookup. + // Prompt one-hot input (multilingual variants only): fill + // prompt_one_hot_in with the resolved language's one-hot replicated + // across T_enc frames (see resolve_prompt_id). if (eb.prompt_one_hot_in != nullptr) { const int P = static_cast(eb.prompt_one_hot_in->ne[0]); const int T_oh = static_cast(eb.prompt_one_hot_in->ne[1]); @@ -1383,35 +1074,21 @@ static transcribe_status run_one_shot_inner( 0, one_hot_buf.size() * sizeof(float)); } - // ----- Host-side sinusoidal positional embedding --------------- + // ----- Host-side sinusoidal positional embedding ----- // - // Sub-stage 3c+: the attention sub-block needs a precomputed - // sin/cos pos_emb input tensor. The encoder graph builder - // allocated `eb.pos_emb_in` with ne=[d_model, pos_len, 1, 1] - // (after it learned T_enc from pre_encode); we fill it here. - // - // Full attention (att_context_left == -1): - // positions[i] = (T_enc - 1) - i for i in [0, 2*T_enc-1) - // Local attention (NeMo LocalAttRelPositionalEncoding): - // positions[i] = W_left - i for i in [0, W_left+W_right+1) - // In both cases the per-row sinusoid is identical: + // Fill eb.pos_emb_in (ne=[d_model, pos_len, 1, 1]). Positions: + // full attention: positions[i] = (T_enc - 1) - i, i in [0, 2T-1) + // local attention: positions[i] = W_left - i, i in [0, W_left+W_right+1) + // Per-row sinusoid (identical in both): // div_term[k] = exp(2k * -ln(10000) / d_model) // pe[i, 2k] = sin(positions[i] * div_term[k]) // pe[i, 2k+1] = cos(positions[i] * div_term[k]) - // - // The on-disk row-major shape is (pos_len, d_model). In ggml ne - // (fast-to-slow) that's [d_model, pos_len, 1, 1] which the - // encoder builder created. if (eb.pos_emb_in != nullptr) { const int d_model = pm->hparams.enc_d_model; const int pos_len = static_cast(eb.pos_emb_in->ne[1]); - // Recover the position of "relative offset 0" inside the buffer. - // - Full attention / chunked_limited streaming: pos_len = 2*T_enc - 1, - // zero index = T_enc - 1 (= (pos_len-1)/2). - // - Regular local attention: pos_len = W_left + W_right + 1, - // zero index = W_left. Only applies to the Regular style — the - // ChunkedLimited path always uses the full RelPositionalEncoding. + // Position of "relative offset 0" in the buffer: (pos_len-1)/2 for + // full / chunked_limited; W_left for Regular local attention. const bool is_chunked = (pm->hparams.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); @@ -1425,7 +1102,6 @@ static transcribe_status run_one_shot_inner( pc->pos_buf.assign(static_cast(pos_len) * d_model, 0.0f); - // div_term[k] = exp(2k * -ln(10000) / d_model) for k in [0, d_model/2) pc->pos_div_term.resize(static_cast(d_model / 2)); const float ln_10000 = std::log(10000.0f); for (int k = 0; k < d_model / 2; ++k) { @@ -1451,13 +1127,10 @@ static transcribe_status run_one_shot_inner( "enc.pos_emb", eb.pos_emb_in, "encoder.pos_emb"); } - // ChunkedLimited attention mask (streaming variants). pos_emb above - // stays at the full 2T-1 length and rel_pos contributes its usual - // (q-k) bias; this mask adds 0 on (q, k) pairs whose chunk indices - // are in [q_chunk - left_chunks, q_chunk] and -INF elsewhere, so - // softmax zeroes everything outside the cache-aware band. NeMo - // semantics: chunk_size = att_context_right + 1, left_chunks = - // att_context_left / chunk_size. + // ChunkedLimited attention mask (streaming variants): 0 on (q, k) + // pairs whose chunk indices are in [q_chunk - left_chunks, q_chunk], + // -INF outside. NeMo: chunk_size = att_context_right + 1, + // left_chunks = att_context_left / chunk_size. if (eb.chunked_mask_in != nullptr) { const int T_enc = static_cast(eb.chunked_mask_in->ne[0]); const int chunk_size = pm->hparams.enc_att_context_right + 1; @@ -1492,14 +1165,10 @@ static transcribe_status run_one_shot_inner( "encoder.attn.chunked_mask"); } - // ----- Set thread count on all backends -------------------------- - // - // Whisper.cpp pattern: iterate the scheduler's backends and set - // n_threads via the registry proc address. GPU backends ignore - // this; CPU and BLAS backends use it. + // Set n_threads on the CPU/BLAS backends (GPU backends ignore it). transcribe::configure_sched_n_threads(pc->sched, pc->n_threads); - // ----- Compute -------------------------------------------------- + // ----- Compute ----- const int64_t t_enc_start = ggml_time_us(); if (const ggml_status gs = ggml_backend_sched_graph_compute(pc->sched, eb.graph); @@ -1513,12 +1182,7 @@ static transcribe_status run_one_shot_inner( pc->t_encode_us = ggml_time_us() - t_enc_start; pc->t_decode_us = 0; - // ----- Dump intermediates (debug only) ------------------------- - // - // Gated on TRANSCRIBE_DUMP_DIR via transcribe::debug::dump_tensor; - // when the env var is unset these are zero-cost. Each field on - // EncoderDumps is a borrowed pointer into the compute_ctx; - // nullptr fields mean "this sub-stage isn't wired yet". + // ----- Dump intermediates (debug only; nullptr fields skipped) ----- auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { @@ -1532,12 +1196,9 @@ static transcribe_status run_one_shot_inner( try_dump("enc.block.0.conv", eb.dumps.block0_after_conv,"encoder.block0.conv"); try_dump("enc.block.0.ff2", eb.dumps.block0_after_ff2, "encoder.block0.ff2"); try_dump("enc.block.0.out", eb.dumps.block0_out, "encoder.block0.out"); - // Mid- and last-block spot-check dumps. File name encodes the - // actual block index (scales with n_layers): 0/12/23 for 24-layer, - // 0/21/41 for 42-layer, 0/8/16 for 17-layer. last_block_out - // aliases final_out (same ggml pointer); we dump it under the - // per-variant block name AND under "enc.final" so both validation - // entries find data. + // Mid- and last-block spot-check dumps. File name encodes the actual + // block index (scales with n_layers). last_block_out aliases + // final_out; dumped under both its block name and "enc.final". if (eb.dumps.mid_block_out != nullptr && eb.dumps.mid_block_idx >= 0) { char name[64]; std::snprintf(name, sizeof(name), "enc.block.%d.out", eb.dumps.mid_block_idx); @@ -1550,13 +1211,8 @@ static transcribe_status run_one_shot_inner( transcribe::debug::dump_tensor(name, eb.dumps.last_block_out, "encoder.block.last.out"); } - // Per-block dump for the layer-by-layer divergence bisect (gated - // by TRANSCRIBE_DUMP_ALL_BLOCKS env var; mark_tensor_for_dump was - // applied per-block in encoder.cpp's loop, so the data is - // preserved across the scheduler's compute. The block 0 / mid / - // last names are still emitted above; this just adds the - // remaining blocks under "enc.block..out" so they collide - // (overwrite) consistently with the named dumps. + // Per-block dump for the layer-by-layer divergence bisect (gated by + // TRANSCRIBE_DUMP_ALL_BLOCKS). if (std::getenv("TRANSCRIBE_DUMP_ALL_BLOCKS") != nullptr) { for (size_t i = 0; i < eb.dumps.all_block_outs.size(); ++i) { ggml_tensor * t = eb.dumps.all_block_outs[i]; @@ -1567,10 +1223,8 @@ static transcribe_status run_one_shot_inner( "encoder.block.bisect"); } } - // Sub-block intermediates for blocks listed in - // TRANSCRIBE_DUMP_SUB_BLOCKS. Each entry was populated by the - // sub-block observer (see encoder.cpp) at graph-build time and - // passed through the scheduler intact via mark_tensor_for_dump. + // Sub-block intermediates for TRANSCRIBE_DUMP_SUB_BLOCKS (populated + // by the sub-block observer in encoder.cpp). for (const auto & p : eb.dumps.sub_block_dumps) { if (p.second == nullptr) continue; transcribe::debug::dump_tensor(p.first.c_str(), p.second, @@ -1579,24 +1233,15 @@ static transcribe_status run_one_shot_inner( try_dump("enc.final", eb.dumps.final_out, "encoder.final"); try_dump("enc.prompted", eb.dumps.prompted_out, "encoder.prompted"); - // Stash the encoder output for the accuracy test (it reaches in - // via the ParakeetSession view) and as a borrowed reference for - // the readback below. pc->encoder_out = eb.out; - // ----- TDT decode -------------------------------------------- - // - // Read the encoder output to host, run the decoder, then build - // the result hierarchy (segments / words / tokens / full text) - // for the public accessors. The encoder activation has ne= - // [d_enc, T_enc, 1, 1] in fast-to-slow ggml order, so the - // contiguous byte range is row-major [T_enc, d_enc] from the - // host's POV — exactly what decode_tdt_greedy expects. + // ----- TDT decode ----- // - // The result snapshot has already been cleared by the caller - // (public run() at its entry, or the streaming dispatcher at - // stream_begin). Re-clearing here would also reset the audio - // cursors stream_finalize is about to commit, so it stays out. + // Read the encoder output to host, decode, build the result + // hierarchy. The activation ne=[d_enc, T_enc, 1, 1] is row-major + // [T_enc, d_enc] from the host's POV — what the decoder expects. The + // caller cleared the result snapshot; re-clearing here would reset + // the cursors stream_finalize is about to commit. const int d_enc = static_cast(eb.out->ne[0]); const int T_enc = static_cast(eb.out->ne[1]); if (d_enc <= 0 || T_enc <= 0) { @@ -1610,13 +1255,9 @@ static transcribe_status run_one_shot_inner( ggml_backend_tensor_get(eb.out, pc->enc_host.data(), 0, pc->enc_host.size() * sizeof(float)); - // For prompt-conditioned models the pre-prompt encoder output is - // a distinct buffer from `eb.out` (which is the post-prompt - // tensor the decoder consumes). The reference dumper labels the - // pre-prompt tensor "dec.enc_out" and the post-prompt tensor - // "dec.enc_out_prompted"; read both back so the validate.py - // compare names match. decode_and_populate handles the prompted - // dump via the override below. + // Prompt-conditioned models: eb.out is the post-prompt tensor + // ("dec.enc_out_prompted"); also read back the pre-prompt final_out + // as "dec.enc_out" so both reference dump names match. if (pm->hparams.has_prompt && eb.dumps.final_out != nullptr && transcribe::debug::enabled()) { @@ -1638,23 +1279,17 @@ static transcribe_status run_one_shot_inner( } // One-shot entry point. Validates session/pm, clears the previous result -// snapshot, then forwards to the shared inference helper. The -// streaming hooks reuse the same helper so a finalize on the -// accumulated buffer produces identical results to a direct run() of -// the same audio. +// snapshot, then forwards to the shared inference helper (also reused by +// the streaming hooks). transcribe_status run( transcribe_session * session, const float * pcm, int n_samples, const transcribe_run_params * params) { - // The dispatcher (transcribe.cpp) has already enum-range - // validated params->task / params->timestamps, rejected - // TRANSLATE against this family's supports_translate=false, and - // rejected any timestamp request finer than our advertised - // max_timestamp_kind=TOKEN. What arrives here is guaranteed - // sane — we only need to resolve AUTO and downcast finer-grained - // output to the requested ceiling when the result is built. + // The dispatcher has already enum-validated params and rejected + // TRANSLATE / sub-TOKEN timestamps; we only resolve AUTO and downcast + // when the result is built. if (session == nullptr || pcm == nullptr || n_samples <= 0) { return TRANSCRIBE_ERR_INVALID_ARG; } @@ -1673,24 +1308,13 @@ transcribe_status run( // Offline batch (transcribe_run_batch) // --------------------------------------------------------------------------- // -// Batches B same-length utterances through ONE encoder graph (n_batch == B, -// the batch riding the activation's ne[2] axis — see the conformer helpers) -// in a single device dispatch, then host-decodes each utterance's encoder -// slice. The mel front-end, decoder, and result-build are identical to the -// single-shot path; only the encoder is fused. Utterances of differing -// length fall back to the per-utterance path (variable-length batching pads -// + masks the overhang — a separate change). Either way every utterance's -// result is captured into session->batch_results in order. - -// Build + compute the batched encoder for `n` utterances, then decode each -// into session->batch_results. Every entry of `mels` is the raw mel-major -// [n_mels, nf[b]] buffer; this packs them into a [T_max, n_mels, 1, n] graph -// input, zero-padding each along time to the batch's T_max. When utterances -// differ in length it also builds + fills the variable-length masks (attn -// key-padding + conv valid-frame) so each utterance's padded tail cannot -// corrupt its real frames, and decodes each at its own T_enc. Same-length -// batches (every nf == T_max) skip the masks and are bit-identical to -// single-shot. +// Batches B utterances through ONE encoder graph (batch on the +// activation's ne[2]) in a single device dispatch, then host-decodes each +// slice; mel / decode / result-build are identical to single-shot. Mels +// pack into [T_max, n_mels, 1, n], zero-padded to T_max. Variable-length +// batches build attn key-padding + conv valid-frame masks so a padded +// tail can't corrupt real frames and decode each at its own T_enc; +// same-length batches skip the masks and are bit-identical to single-shot. static int pre_encode_t_out(int in) { // One stride-2, kernel-3, pad-1 conv: floor((in + 2 - 3)/2) + 1. return (in - 1) / 2 + 1; @@ -1835,10 +1459,9 @@ static transcribe_status run_batch_encode( 0, pc->pos_buf.size() * sizeof(float)); } - // Prompt one-hot upload (multilingual variants). Resolves params->language - // ONCE for the whole batch (the public ABI doesn't expose a per-utterance - // language array yet — every utterance in a batched call shares the same - // language hint) and replicates the one-hot across all (T_enc, B) slots. + // Prompt one-hot upload (multilingual variants). Resolves + // params->language once for the whole batch (no per-utterance language + // array in the ABI) and replicates across all (T_enc, B) slots. if (eb.prompt_one_hot_in != nullptr) { const int P = static_cast(eb.prompt_one_hot_in->ne[0]); const int T_oh = static_cast(eb.prompt_one_hot_in->ne[1]); @@ -1866,12 +1489,9 @@ static transcribe_status run_batch_encode( 0, one_hot_buf.size() * sizeof(float)); } - // ChunkedLimited attention mask (cache-aware variants). Allocated by - // build_encoder_graph when the model declares chunked attention; the - // single-shot path fills it and the batched path must too (otherwise the - // input is uninitialized and corrupts every attention score). The pattern - // depends only on T_enc, which the whole batch shares, so one mask - // broadcasts across the batch. Mirrors run_one_shot_inner. + // ChunkedLimited attention mask (cache-aware variants), same as + // run_one_shot_inner. The pattern depends only on T_enc, shared + // across the batch, so one mask broadcasts. if (eb.chunked_mask_in != nullptr) { const int Tk = static_cast(eb.chunked_mask_in->ne[0]); const int chunk_size = pm->hparams.enc_att_context_right + 1; @@ -1932,11 +1552,9 @@ static transcribe_status run_batch_encode( ggml_backend_tensor_get(eb.out, pc->enc_host.data(), 0, pc->enc_host.size() * sizeof(float)); - // Prompt-conditioned variants: eb.out is the POST-prompt tensor (what - // the decoder consumes); the single-shot path dumps the PRE-prompt - // tensor as "dec.enc_out" and the POST-prompt as "dec.enc_out_prompted" - // so the comparator sees both under their reference names. Mirror that - // in the batched path with per-utterance .b{i} suffixes. + // Prompt-conditioned variants: dump the PRE-prompt final_out per + // utterance as "dec.enc_out.b{i}" (the post-prompt eb.out is dumped + // via decode_and_populate's override). Mirrors single-shot. const bool has_prompt = pm->hparams.has_prompt && eb.dumps.final_out != nullptr; std::vector unprompted_host; @@ -1959,10 +1577,9 @@ static transcribe_status run_batch_encode( const char * enc_dump_name = has_prompt ? "dec.enc_out_prompted" : nullptr; - // Host-slice the shared encoder output and decode each utterance. The - // encoder is ONE shared dispatch, so decode_batch_slices amortizes its - // total compute + the total mel cost across the batch (the per-utt sum then - // equals the real batch time); decode is genuinely per-utterance. + // Host-slice the shared encoder output and decode each utterance. + // decode_batch_slices amortizes the one shared encoder + total mel + // cost across the batch; decode is genuinely per-utterance. return transcribe::decode_batch_slices( pc, n, pc->enc_host.data(), utt_elems, pc->t_encode_us, total_mel_us, [&](int b, const float * enc_b) { @@ -1989,15 +1606,10 @@ transcribe_status run_batch( } transcribe::debug::init(); - // Compute each utterance's mel. A malformed utterance (null pcm / - // non-positive samples / mel failure) means we cannot pack a clean batch - // tensor, so we fall back to the per-utterance path for the whole call - // (rare; keeps the batch tensor rectangular and the masks well-defined). - // The mel front-end is pure host code with no cross-utterance state, so - // the B extractions run in parallel across CPU workers (see - // transcribe::parallel_for_all) — frequently the dominant wall cost once - // the encoder is on a fast accelerator. n_mels is constant across - // utterances; collect it per-index to avoid a shared write. + // Compute each utterance's mel in parallel (pure host, no + // cross-utterance state). A malformed utterance falls the whole call + // back to the per-utterance path (keeps the batch tensor rectangular). + // n_mels collected per-index to avoid a shared write. std::vector> mels(static_cast(n)); std::vector nf(static_cast(n), 0); std::vector n_mels_per(static_cast(n), 0); @@ -2044,25 +1656,14 @@ transcribe_status run_batch( } // --------------------------------------------------------------------------- -// Streaming hooks (Phase 4a: stream-of-whole) +// Streaming hooks // --------------------------------------------------------------------------- // -// Phase 4a buffers the entire stream in stream_pcm_buffer and runs -// the existing one-shot inference path at finalize. This makes the -// streaming API observable end-to-end and trivially parity-equivalent -// to transcribe_run on the same audio. Real incremental encoder / -// RNNT-predictor streaming via NeMo's cache_last_channel / -// cache_last_time tensors is Phase 4b. -// -// Only cache-aware streaming variants (ChunkedLimited attention) take -// this path; offline parakeet variants never reach the hook because -// load() leaves caps.supports_streaming = false for them and the -// streaming dispatcher rejects begin with NOT_IMPLEMENTED. The defensive -// gate in stream_begin below is defense in depth. -// -// The dispatcher handles the IDLE/ACTIVE/FINISHED/FAILED lifecycle -// and the result-snapshot counters (revision, committed counts, audio -// cursors). Family hooks own the per-utterance audio scratch. +// Only streaming variants reach these hooks (load() leaves +// supports_streaming = false otherwise; the gate in stream_begin is +// defense in depth). The dispatcher handles the lifecycle and the +// result-snapshot counters; the family hooks own the per-utterance audio +// scratch. constexpr int64_t k_sample_rate_hz = 16000; @@ -2075,34 +1676,17 @@ static int64_t us_to_ms(int64_t us) { } // --------------------------------------------------------------------------- -// M2 streaming-encoder helpers +// Streaming-encoder helpers // --------------------------------------------------------------------------- - -// Per-chunk shape constants for nemotron-speech-streaming-en-0.6b -// (att_context_size=[70,13], subsampling_factor=8). These match NeMo's -// Streaming chunk geometry now lives on ParakeetStreamingCaches and is -// resolved per stream from ParakeetHParams + the caller-selected -// att_context_right. The values used here in stream_feed / -// emit_streaming_chunk are pc->stream_caches.{mel_chunk_total, -// mel_new_per_chunk, drop_extra_pre_encoded}; pm->hparams. -// enc_stream_pre_encode_cache_size is used directly for the -// history-frame count (= mel_chunk_total - mel_new_per_chunk). // -// Simplification vs NeMo (M2): we always use the "subsequent" mode -// (with history-frame mel prepend) for every chunk including the -// first; the first chunk's history is zero-initialized at -// stream_begin. The downside is a few frames of output are wasted at -// the start; the upside is uniform driver logic. Phase B will fix -// the first-chunk semantics to match NeMo's chunk_size[0] / -// pre_encode_cache_size[0] = 0. - -// Fill the sinusoidal pos_emb buffer for a streaming chunk. Layout is -// identical to the offline ChunkedLimited path (RelPositionalEncoding): -// pos_len = 2*T_virtual - 1; positions[i] = (T_virtual - 1) - i; per-row -// sinusoid uses div_term[k] = exp(2k * -ln(10000) / d_model). -// `pos_emb_in` shape: ne=[d_model, pos_len, 1, 1] f32 (the encoder -// graph builder allocated it). Side-effect: pc->pos_buf and -// pc->pos_div_term are resized in place. +// Chunk geometry lives on ParakeetStreamingCaches, resolved per stream +// from ParakeetHParams + the caller-selected att_context_right. + +// Fill the sinusoidal pos_emb buffer for a streaming chunk. Same layout +// as the offline ChunkedLimited path: pos_len = 2*T_virtual - 1, +// positions[i] = (T_virtual - 1) - i, div_term[k] = exp(2k * -ln(10000) +// / d_model). pos_emb_in is ne=[d_model, pos_len, 1, 1] f32. Resizes +// pc->pos_buf / pc->pos_div_term in place. void fill_streaming_pos_emb(ParakeetSession * pc, ggml_tensor * pos_emb_in, int d_model, @@ -2178,13 +1762,10 @@ void fill_streaming_chunked_mask(ggml_tensor * mask_in, } // Rebuild the per-layer rel-pos projection cache for `pos_len` (see -// ParakeetStreamingCaches::pos_proj). Runs a small one-off graph — -// n_layers mul_mats of [d_model, d_model] x [d_model, pos_len] plus -// layout massaging — through the session scheduler, reading the pos_emb -// values from pc->pos_buf (filled by fill_streaming_pos_emb earlier in -// the same chunk). Called only on a geometry change: in practice twice -// per stream session (first-chunk geometry, then steady state) and -// never again, since the cache is geometry-keyed, not stream-keyed. +// ParakeetStreamingCaches::pos_proj). Runs a one-off graph (n_layers +// mul_mats + layout massaging) through the session scheduler, reading +// pos_emb from pc->pos_buf. Called only on a geometry change — twice per +// stream in practice (first-chunk, then steady state). transcribe_status ensure_pos_proj_cache(ParakeetSession * pc, ParakeetModel * pm, int pos_len) @@ -2284,9 +1865,8 @@ transcribe_status ensure_pos_proj_cache(ParakeetSession * pc, } // namespace -// Already inside namespace transcribe::parakeet { ... }. The function is -// declared in parakeet.h so external callers (unit tests, the buffered -// streaming driver) can link against it. +// Declared in parakeet.h (non-anonymous) so unit tests and the buffered +// driver can link against it. void compute_chunked_limited_with_rc_mask( float * out_buf, int T, @@ -2333,25 +1913,15 @@ namespace { // Build, run, and post-process a single streaming encoder chunk. // -// Inputs: -// pc, pm context + model -// mel_chunk_data row-major [n_mels, n_mel_chunk_frames] -// f32 (matches ggml ne=[n_mel_chunk_frames, -// n_mels, 1, 1] for upload). -// n_mel_chunk_frames mel frames in this chunk (e.g. 121 for -// nemotron subsequent chunks = 9 cache + -// 112 new; 105 for first chunk = no -// cache prepend). -// drop_extra_pre_encoded 0 for first chunk, 2 for subsequent on -// nemotron-streaming. -// mel_frames_advance how many mel frames this chunk consumes -// from the input buffer (chunk_size_first -// or chunk_size_subsequent). NOT the same -// as n_mel_chunk_frames when a pre-encode -// cache prepend is in play. -// -// On success, appends new TdtTokens to pc->raw_tokens and advances -// the persistent cache + decoder state. +// mel_chunk_data row-major [n_mels, n_mel_chunk_frames] f32. +// n_mel_chunk_frames mel frames in this chunk (history prepend + +// new). +// drop_extra_pre_encoded 0 for first chunk, else drop_extra_subsequent. +// mel_frames_advance mel frames this chunk consumes from the input +// (NOT n_mel_chunk_frames when a cache prepend +// is in play). +// On success, appends TdtTokens to pc->raw_tokens and advances the cache +// + decoder state. transcribe_status emit_streaming_chunk( ParakeetSession * pc, ParakeetModel * pm, @@ -2367,14 +1937,13 @@ transcribe_status emit_streaming_chunk( if (pc->kv_type == TRANSCRIBE_KV_TYPE_F32) resolved_kv = GGML_TYPE_F32; if (pc->kv_type == TRANSCRIBE_KV_TYPE_F16) resolved_kv = GGML_TYPE_F16; - // Tear down any previous compute_ctx (one per chunk; matches the - // offline run() lifecycle and keeps scratch bounded). + // One compute_ctx per chunk (matches the offline run() lifecycle). if (pc->compute_ctx != nullptr) { ggml_free(pc->compute_ctx); pc->compute_ctx = nullptr; } ggml_init_params ip {}; - ip.mem_size = 16 * 1024 * 1024; // 16 MB; same as offline run path + ip.mem_size = 16 * 1024 * 1024; ip.mem_buffer = nullptr; ip.no_alloc = true; pc->compute_ctx = ggml_init(ip); @@ -2465,35 +2034,25 @@ transcribe_status emit_streaming_chunk( const int T_cache = hp.enc_att_context_left; const int T_q_new = T_virtual - T_cache; - // Build & upload pos_emb. The zero-offset row sits at T_virtual - 1 - // in both the square (pos_len = 2*T_virtual-1) and the query-sliced - // rectangular (pos_len = T_virtual + T_q_new - 1) geometries. Null - // when every block consumes the memoized rel-pos projection — the - // graph then carries no pos_emb input at all. + // Build & upload pos_emb. The zero-offset row sits at T_virtual - 1. + // Null when every block consumes the memoized rel-pos projection. if (eb.pos_emb_in != nullptr) { fill_streaming_pos_emb(pc, eb.pos_emb_in, hp.enc_d_model, /*zero_index=*/T_virtual - 1); } - // Build & upload chunked mask with cache-unfilled prefix masking. - // The mask band is a function of the (att_context_left, att_context_right) - // RESOLVED for this stream from the caller's latency selection, not the - // model-default hparams. They coincide today (left is invariant across - // the menu, and the per-chunk geometry derived from chosen_right keeps - // the default-R chunking accidentally equivalent on the new-frame rows), - // but reading the resolved values makes the mask correct by construction - // for any future variant whose menu breaks those invariants. + // Build & upload the chunked mask with cache-unfilled prefix masking. + // The band uses the (left, right) RESOLVED for this stream, not the + // model-default hparams (correct by construction even if a future + // variant's menu breaks today's coincidence). fill_streaming_chunked_mask( eb.chunked_mask_in, T_virtual, T_cache, pc->stream_caches.channel_len, pc->stream_caches.att_context_left, pc->stream_caches.att_context_right); - // Prompt one-hot upload (multilingual variants only). Same shape - // contract as the offline path: [num_prompts, T_q_new, 1, 1] holds - // a single one-hot column at the resolved language's index, - // replicated across T_q_new. The streaming language hint lives in - // pc->stream_run_params (captured at stream_begin). + // Prompt one-hot upload (multilingual variants only). Same shape as + // the offline path; the language hint lives in pc->stream_run_params. if (eb.prompt_one_hot_in != nullptr) { const int P = static_cast(eb.prompt_one_hot_in->ne[0]); const int T_oh = static_cast(eb.prompt_one_hot_in->ne[1]); @@ -2541,19 +2100,14 @@ transcribe_status emit_streaming_chunk( ggml_backend_tensor_get(eb.out, pc->enc_host.data(), 0, pc->enc_host.size() * sizeof(float)); - // Dump enc_out + cache_out (snapshot of outputs from this chunk). - // Done BEFORE the cache rotation so we read the freshly-computed - // cache_out tensors. channel_len is dumped post-rotation since - // that matches NeMo's "cache_last_channel_len after chunk" return. + // Dump enc_out + cache_out, BEFORE the cache rotation so the + // cache_out tensors are still the freshly-computed ones. if (dump_on) { char namebuf[128]; std::snprintf(namebuf, sizeof(namebuf), "stream.chunk.%d.enc_out", step_num); - // Match the Python dumper's to_np() squeeze: when T_q_new == 1 - // (R=0 at the lowest-latency setting), drop the leading size-1 - // dim so the on-disk shape becomes [d_enc] not [1, d_enc]. - // compare_tensors.py treats shape mismatches as failures - // regardless of element values. + // Match the Python dumper's squeeze: T_q_new == 1 (R=0) drops the + // leading size-1 dim so the on-disk shape is [d_enc] not [1, d_enc]. if (T_q_new == 1) { const long long enc_shape[1] = {d_enc}; transcribe::debug::dump_host_f32( @@ -2580,16 +2134,11 @@ transcribe_status emit_streaming_chunk( } } - // Rotate per-layer caches: cache_out → persistent cache_in. Both - // tensors live on the model's primary backend, so a backend-side - // copy avoids host roundtripping. The null check is defense in - // depth — every supported R produces a fully-allocated cache_out - // (the conv_module cache_next logic in conformer/conformer.cpp - // covers T_q_new < pad_left explicitly), so reaching this branch - // indicates a builder bug. - // KV-cache mode rotates last_k/last_v and leaves the channel cache - // untouched (the builder set channel_out to null); the recompute - // path rotates last_channel. The time (conv) cache rotates in both. + // Rotate per-layer caches: cache_out → persistent cache_in + // (backend-side copy, no host roundtrip). KV-cache mode rotates + // last_k/last_v (channel_out is null); the recompute path rotates + // last_channel. The time (conv) cache rotates in both. An unallocated + // cache_out here is a builder bug. const bool kv_mode = !cache_io.k_out.empty(); for (int i = 0; i < n_layers; ++i) { ggml_tensor * ch = cache_io.channel_out[i]; @@ -2618,13 +2167,10 @@ transcribe_status emit_streaming_chunk( pc->stream_caches.channel_len = std::min( T_cache, pc->stream_caches.channel_len + T_q_new); - // Refill the rel-pos projection memo on geometry change, so the next - // chunk with this pos_len consumes the cached projections instead of - // running the per-layer linear_pos matmuls. Must come after the cache - // rotation above: ensure_pos_proj_cache resets the scheduler, which - // releases this chunk graph's compute buffers (cache_out lives there). - // A failure only loses the memoization — streaming continues on the - // inline-compute path — so it logs and degrades. + // Refill the rel-pos projection memo on geometry change. Must come + // after the cache rotation: ensure_pos_proj_cache resets the + // scheduler, releasing this chunk's compute buffers (cache_out lives + // there). A failure only loses the memoization, so it logs and degrades. if (eb.pos_emb_in != nullptr) { const int pos_len_cur = static_cast(eb.pos_emb_in->ne[1]); if (pos_len_cur != pc->stream_caches.pos_proj_len) { @@ -2666,28 +2212,17 @@ transcribe_status emit_streaming_chunk( } pc->stream_dec_state.frame_offset += T_q_new; - // Cursor advances by mel_frames_advance (caller computes: - // chunk_size_first on the first chunk, chunk_size_subsequent - // afterward). This is the same as NeMo's shift_size: it is - // INDEPENDENT of the chunk we fed (which may include a - // pre-encode-cache prepend). + // Cursor advances by mel_frames_advance (NeMo's shift_size), + // INDEPENDENT of the fed chunk (which may include a cache prepend). pc->stream_caches.mel_frames_consumed += mel_frames_advance; return TRANSCRIBE_OK; } // Rebuild the public result vectors (tokens, full_text) from the -// committed raw_tokens. Simple and idempotent: clears the vectors and -// re-populates from scratch every call. Segments / words are NOT -// populated during streaming — those derive from richer logic in the -// offline run() result builder and are deferred to stream_finalize -// (or omitted entirely for the first M2 cut). -// -// result_kind is TOKEN whenever any tokens are present: each TokenEntry -// gets real t0_ms / t1_ms from step_at_emit, so the snapshot is honest -// at token granularity. Words/segments stay empty until finalize, but -// a caller that asked for WORD or SEGMENT timestamps already passed the -// max_timestamp_kind gate at begin time (Parakeet advertises TOKEN as -// its max), so TOKEN here is the strongest honest answer we can give. +// committed raw_tokens; idempotent. Words/segments are NOT populated +// during streaming (deferred to finalize). result_kind is TOKEN whenever +// tokens are present — each gets real t0/t1 from step_at_emit, the +// strongest honest answer at this granularity. void rebuild_streaming_result_text(ParakeetSession * pc, const ParakeetModel * pm) { @@ -2702,19 +2237,11 @@ void rebuild_streaming_result_text(ParakeetSession * pc, return; } - // Frame-to-ms ratio, identical to the offline builder and NeMo's - // reference (see parakeet_ms_per_enc_frame): float multiply, rounded once - // per token below. const double frame_to_ms = parakeet_ms_per_enc_frame(pm->hparams); - // Strip multilingual language tags (e.g. the trailing "" - // / "" the nemotron-3.5 vocab emits) and other CONTROL pieces - // from the public streaming result by default, mirroring the offline - // builder (decode_and_populate) so streaming honors the same documented - // contract. Gated on the keep_special_tags run param captured at - // stream_begin (CLI --raw-tokens). Only the public projection is - // filtered: pc->raw_tokens stays whole, so the decoder/commit cursor and - // the raw token accessors are unaffected. + // Strip CONTROL / tag pieces from the public result, mirroring + // decode_and_populate (gated on keep_special_tags). Only the public + // projection is filtered; pc->raw_tokens stays whole. const transcribe::Tokenizer & tok = pm->tok; const bool strip_tags = !pc->stream_run_params.keep_special_tags; @@ -2745,30 +2272,16 @@ void rebuild_streaming_result_text(ParakeetSession * pc, pc->result_kind = TRANSCRIBE_TIMESTAMPS_TOKEN; } -// ----- Buffered streaming (parakeet-unified-en-0.6b) ----------------- +// ----- Buffered streaming (parakeet-unified-en-0.6b) ----- // -// Mirrors NeMo's `speech_to_text_streaming_infer_rnnt.py` reference at -// the per-chunk granularity. The variable-stride algorithm: -// - Step 0 (initial fill): num_new = samples_chunk + samples_right. -// - Steady state: num_new = samples_chunk. -// - Final step (from finalize): num_new = remaining audio; the -// trailing right context slot is folded into chunk via -// `add_frames_get_removed_(is_last=true)`. -// Per step: -// - Update the buffer's internal ContextSize (buf_ctx_*) via the -// same accounting NeMo's StreamingBatchedAudioBuffer uses. -// - Slice the last session.total() samples from stream_pcm_buffer as -// the encoder window. -// - Compute mel over the full window. -// - Build the encoder graph with a BufferedStreamMaskOverride to -// engage the chunked_limited_with_rc attention mask sized for -// the expected (L, C, R). -// - Slice off the first (ctx_left / samples_per_frame) encoder -// frames and decode (ctx_chunk / samples_per_frame) on non-last -// or (T_enc_full - ctx_left_frames) on last with greedy RNN-T -// plus carried LstmState. Matches the reference's -// encoder_context_batch.chunk vs -// encoder_output_len - encoder_context_batch.left dispatch. +// Mirrors NeMo's speech_to_text_streaming_infer_rnnt.py. Variable-stride +// per step: step 0 num_new = samples_chunk + samples_right; steady state +// num_new = samples_chunk; the final step (finalize) consumes the rest +// and folds the right slot into chunk. Each step updates the buffer's +// ContextSize (buf_ctx_*), slices the encoder window, computes mel, +// builds the graph with a BufferedStreamMaskOverride, then slices off the +// ctx_left frames and decodes ctx_chunk (or all remaining on last) with a +// carried RNN-T LstmState. static void buf_ctx_add_frames( ParakeetSession * pc, int64_t num_new, bool is_last) { @@ -2831,12 +2344,8 @@ transcribe_status emit_buffered_chunk( } } - // ----- Per-chunk dumps (TRANSCRIBE_DUMP_DIR; mirror NeMo names) ----- - // - // Push a `stream.chunk..` prefix so every dump_tensor inside the - // encoder graph builder gets scoped per chunk (otherwise the block - // outputs `enc.block..out` would overwrite across chunks). The - // prefix is popped at function exit via the guard below. + // Push a `stream.chunk..` dump prefix so per-chunk block outputs + // don't overwrite across chunks; popped at function exit. const int chunk_step = pc->buf_chunk_step; char chunk_prefix[64]; std::snprintf(chunk_prefix, sizeof(chunk_prefix), @@ -2949,13 +2458,9 @@ transcribe_status emit_buffered_chunk( 0, pc->pos_buf.size() * sizeof(float)); } - // ----- Conformer conv pad mask ----- - // - // NeMo passes pad_mask into every Conformer conv module and zeros - // post-GLU activations where pre_encode emitted padded overhang - // frames. The attention mask below already excludes those frames - // from MHA; this mask prevents them from leaking backward through - // the depthwise conv's right context. + // Conv pad mask: zeros post-GLU activations on pre_encode overhang + // frames so they don't leak backward through the depthwise conv's + // right context (the attention mask already excludes them from MHA). if (eb.conv_pad_mask_in != nullptr) { const int T_enc = static_cast(eb.conv_pad_mask_in->ne[0]); @@ -2972,19 +2477,11 @@ transcribe_status emit_buffered_chunk( "encoder.conv.pad_mask"); } - // ----- Chunked_limited_with_rc mask fill ----- - // - // Pad-mask: NeMo's encoder ANDs a conv-overhang pad_mask onto the - // chunked mask before MHA — frames past `padding_length` (post - // pre-encode) are masked out. The conv subsampling can emit T_enc - // > session.total/samples_per_frame when the input doesn't tile the - // subsampling stride exactly; without the pad mask, those trailing - // frames contaminate every other frame's attention scores. We - // mirror NeMo by passing `effective_T = session.total / - // samples_per_frame` so the mask compute folds in the pad_mask. - // Critical at low-C/low-R configs where one extra frame of - // contamination tips many greedy-emission decisions; barely - // noticeable at (70,13,13). + // Chunked_limited_with_rc mask. effective_T folds in NeMo's + // conv-overhang pad_mask: the subsampling can emit T_enc > + // session.total/samples_per_frame and those trailing frames would + // otherwise contaminate every frame's attention scores (critical at + // low-C/low-R). if (eb.chunked_mask_in != nullptr) { const int T_enc = static_cast(eb.chunked_mask_in->ne[0]); @@ -3018,20 +2515,13 @@ transcribe_status emit_buffered_chunk( pc->encoder_out = eb.out; - // Dump cpp's mel input alongside the per-block intermediates so the - // ref-vs-cpp bisect can also localize drift to the mel frontend. if (transcribe::debug::enabled()) { transcribe::debug::dump_tensor("enc.mel.in", eb.mel_in, "encoder.mel"); } // ----- Per-chunk intermediate dumps (debug only) ----- - // - // Same set of intermediates the offline path dumps via run(); used - // for layer-by-layer divergence bisect against NeMo's per-block - // outputs at the streaming geometry. Gated on TRANSCRIBE_DUMP_DIR - // and TRANSCRIBE_DUMP_ALL_BLOCKS env vars. The active per-chunk - // name prefix (push_name_prefix above) scopes these names to this - // chunk's directory: `stream.chunk..enc.block..out` etc. + // Same set as the offline run() path, scoped by the active per-chunk + // name prefix. if (transcribe::debug::enabled()) { auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) { @@ -3100,9 +2590,7 @@ transcribe_status emit_buffered_chunk( }; if (T_enc_full <= ctx_left_frames) { - // Window was too small to produce any chunk-region frames - // (e.g. degenerate audio shorter than the left context). Skip - // emission for this chunk. + // Window too small for any chunk-region frames; skip emission. advance_cursor(); return TRANSCRIBE_OK; } @@ -3122,12 +2610,9 @@ transcribe_status emit_buffered_chunk( const float * enc_chunk = pc->enc_host.data() + static_cast(ctx_left_frames) * static_cast(d_enc); - // Per-chunk encoder output dump. Mirrors NeMo's reference: emit - // the FULL post-slice encoder output (chunk + right context), - // not just the frames that get decoded this step. The decoder - // gates on T_to_decode below; the dump captures the entire - // chunk-aligned region so the parity harness can compare both - // the decoded and lookahead portions. + // Per-chunk encoder output dump: the FULL post-slice output (chunk + + // right context), not just the T_to_decode frames, so the parity + // harness sees both the decoded and lookahead portions. if (transcribe::debug::enabled()) { const long long shape[2] = { static_cast(T_chunk_avail), @@ -3159,11 +2644,9 @@ transcribe_status emit_buffered_chunk( return TRANSCRIBE_OK; } -// Pure validation of the buffered-stream extension. Resolves (L, C, R) -// to encoder frames using the same sentinel semantics stream_begin uses -// to configure the buffered path; emits the same stderr diagnostics on -// rejection. Does NOT touch pc->buf_*. Output pointers are only written -// on TRANSCRIBE_OK return. +// Pure validation of the buffered-stream extension: resolves (L, C, R) +// to encoder frames. Does NOT touch pc->buf_*; out pointers written only +// on TRANSCRIBE_OK. transcribe_status resolve_buffered_stream_geom( const ParakeetModel * pm, const transcribe_stream_params * stream_params, @@ -3252,10 +2735,9 @@ transcribe_status resolve_buffered_stream_geom( return TRANSCRIBE_OK; } -// Pure validation of the cache-aware-stream extension. Resolves the -// active (att_context_left, att_context_right) for the stream from the -// model's training menu. Does NOT touch pc->*. Output pointers are only -// written on TRANSCRIBE_OK return. +// Pure validation of the cache-aware-stream extension: resolves the +// active (att_context_left, att_context_right) from the training menu. +// Does NOT touch pc->*; out pointers written only on TRANSCRIBE_OK. transcribe_status resolve_cache_aware_stream_geom( const ParakeetModel * pm, const transcribe_stream_params * stream_params, @@ -3314,11 +2796,9 @@ transcribe_status resolve_cache_aware_stream_geom( return TRANSCRIBE_OK; } -// Pre-flight: validate caller-supplied extension fields without mutating -// session or model state. Called by the dispatcher BEFORE clear_result, -// so a rejection here leaves the previous utterance's snapshot intact. -// stream_begin re-runs the same resolvers as defense in depth and uses -// their parsed outputs to configure per-stream state. +// Pre-flight: validate caller extension fields without mutating state. +// Called by the dispatcher before clear_result, so a rejection leaves the +// previous snapshot intact. stream_begin re-runs the same resolvers. transcribe_status stream_validate( const transcribe_session * session, const transcribe_run_params * /*run_params*/, @@ -3360,15 +2840,10 @@ transcribe_status stream_begin( return TRANSCRIBE_ERR_INVALID_ARG; } - // Initialize the debug dumper from TRANSCRIBE_DUMP_DIR. Idempotent. - // Streaming may run without ever going through run() (the offline - // path), so this is needed in addition to the run() callsite. + // Streaming may run without going through run(), so init the dumper here. transcribe::debug::init(); - // Defense in depth: the dispatcher should have already rejected - // begin via the supports_streaming gate for non-streaming variants, - // but reject again here in case a future refactor lets a Regular - // (offline) variant slip through. + // Defense in depth (the dispatcher already gates on supports_streaming). const bool is_chunked_limited = (pm->hparams.enc_att_context_style == ParakeetHParams::AttContextStyle::ChunkedLimited); @@ -3381,18 +2856,13 @@ transcribe_status stream_begin( // -------- Buffered streaming path (parakeet-unified-en-0.6b) -------- // - // The unified model declares chunked_limited_with_rc and ships a - // 3-tuple training menu. Buffered streaming re-runs the offline - // encoder on a sliding [left | chunk | right] PCM window per chunk; - // there's no per-layer cache to maintain. The RNN-T predictor + - // joint reuse the existing greedy-streaming decoder so the only - // new state is the runtime (L, C, R) geometry and the LSTM carry. + // chunked_limited_with_rc with a 3-tuple training menu. Re-runs the + // offline encoder on a sliding [left | chunk | right] PCM window per + // chunk (no per-layer cache); the only new state is the runtime + // (L, C, R) geometry and the LSTM carry. if (is_chunked_with_rc) { - // Resolve and validate (L, C, R) in encoder frames. The same - // helper runs in stream_validate, so the dispatcher has - // already rejected bad caller input before clearing the - // previous result; the call here is defense in depth and - // also produces the parsed values we configure state from. + // Resolve (L, C, R) in encoder frames (defense in depth; also + // produces the parsed values we configure state from). int L_frames = 0, C_frames = 0, R_frames = 0; if (const transcribe_status st = resolve_buffered_stream_geom( pm, stream_params, &L_frames, &C_frames, &R_frames); @@ -3423,9 +2893,7 @@ transcribe_status stream_begin( pc->raw_tokens.clear(); pc->stream_run_params = *run_params; reset_streaming_decoder_state(pc, pm); - // For buffered streaming the predictor state's frame_offset - // starts at 0; we advance it per-chunk by T_to_decode so token - // step_at_emit lands in stream-wide encoder-frame coordinates. + // frame_offset starts at 0, advanced per-chunk by T_to_decode. pc->stream_dec_state.frame_offset = 0; return TRANSCRIBE_OK; @@ -3434,10 +2902,8 @@ transcribe_status stream_begin( // -------- Cache-aware streaming path (nemotron-speech-streaming) -------- pc->buf_active = false; - // Resolve the active (att_context_left, att_context_right) for this - // stream. Same helper runs in stream_validate; rerunning here is - // defense in depth and also extracts the values we use to configure - // the per-stream caches. + // Resolve the active (att_context_left, att_context_right) (defense in + // depth; also extracts the values we configure the caches from). int chosen_left = 0, chosen_right = 0; if (const transcribe_status st = resolve_cache_aware_stream_geom( pm, stream_params, &chosen_left, &chosen_right); @@ -3447,19 +2913,13 @@ transcribe_status stream_begin( } pc->stream_pcm_buffer.clear(); - pc->raw_tokens.clear(); // The dispatcher clears session->tokens / words / - // segments / full_text via clear_result, but - // raw_tokens is parakeet-internal and - // accumulates per-chunk; without this reset - // a back-to-back stream_begin (batch WER) on - // the same context would carry the previous - // utterance's tokens into the next stream. + pc->raw_tokens.clear(); // parakeet-internal, accumulates per-chunk; + // clear_result doesn't touch it, so reset here + // or a back-to-back stream_begin carries tokens. pc->stream_run_params = *run_params; - // M2: allocate streaming caches on first stream_begin for this - // context (idempotent — survives across utterances). Zero the - // contents and reset cursors on every begin so each stream starts - // from NeMo's get_initial_cache_state (zeros). + // Allocate streaming caches on first stream_begin (idempotent); zero + // contents and reset cursors on every begin. if (const transcribe_status st = init_streaming_caches(pc, pm); st != TRANSCRIBE_OK) { @@ -3468,26 +2928,16 @@ transcribe_status stream_begin( zero_streaming_caches(pc); reset_streaming_decoder_state(pc, pm); - // Resolve chunking geometry for this stream. Matches NeMo's - // setup_streaming_params formulas exactly: - // + // Resolve chunking geometry (NeMo's setup_streaming_params): // chunk_size_subsequent = subsampling_factor * (1 + R) // chunk_size_first = sampling_frames_first + subsampling_factor * R - // = chunk_size_subsequent - (subsampling_factor - // - sampling_frames_first) - // shift_size_subsequent = chunk_size_subsequent (cache_drop_size=0 - // for chunked_limited) - // mel_fed_first = chunk_size_first (pre_encode_cache_size[0] = 0) + // mel_fed_first = chunk_size_first // mel_fed_subsequent = chunk_size_subsequent + pre_encode_cache_size // drop_extra_first = 0 // drop_extra_subsequent = streaming_cfg.drop_extra_pre_encoded - // - // Fallbacks for legacy streaming GGUFs (converted before the new - // streaming.* KVs landed): derive pre_encode_cache_size from - // subsampling_factor+1, drop_extra_pre_encoded via NeMo's formula, - // and sampling_frames_first defaults to subsampling_factor (which - // collapses chunk_size_first == chunk_size_subsequent — i.e. no - // first-chunk special case, matching pre-Phase-B behavior). + // Fallbacks below cover legacy GGUFs lacking the streaming.* KVs + // (sampling_frames_first defaults to subsampling_factor, collapsing + // chunk_size_first == chunk_size_subsequent — no first-chunk case). const int subsampling_factor = pm->hparams.enc_subsampling_factor; int pre_encode_cache_size = pm->hparams.enc_stream_pre_encode_cache_size; if (pre_encode_cache_size <= 0) { @@ -3545,15 +2995,11 @@ transcribe_status stream_feed( // -------- Buffered streaming path -------- // - // Emit one chunk at a time while we have enough audio in - // stream_pcm_buffer to cover the next non-final [chunk | right] - // window. The left context is pulled from already-buffered audio - // (zero-padded at the start of the stream). + // Emit non-last chunks while there's enough buffered audio for the + // next [chunk | right] window (the last chunk is stream_finalize's + // job). Step 0 needs samples_chunk + samples_right; steady-state + // needs samples_chunk. if (pc->buf_active) { - // Variable-stride loop: step 0 requires (samples_chunk + - // samples_right) of new audio; steady-state needs samples_chunk. - // We only emit non-last chunks here; the last chunk (and its - // ragged-tail folding) is the responsibility of stream_finalize. while (true) { if (pc->poll_abort()) return TRANSCRIBE_ERR_ABORTED; const int64_t num_new = pc->buf_initialized @@ -3596,20 +3042,12 @@ transcribe_status stream_feed( // -------- Cache-aware streaming path -------- // - // Recompute mel from the (sliding) accumulated PCM buffer. - // - // The buffer is trimmed after each emit to retain only the prefix - // samples that future mel frames will still need (the STFT window - // of the next-unemitted frame extends back by `n_fft/2` samples). - // That keeps mel cost bounded per feed regardless of stream length. - // - // pcm_start_sample is the absolute sample index of buffer[0] and - // is always hop-aligned, so the absolute → buffer-relative - // mel-frame translation is the integer pcm_start_sample / hop. - // The first few buffer-relative mel frames (where the STFT window - // would extend left of buffer[0]) are reflect-pad-contaminated - // garbage; they correspond to already-emitted absolute frames and - // are never read. + // Recompute mel from the sliding PCM buffer (trimmed after each emit + // to keep per-feed mel cost bounded). pcm_start_sample is the absolute + // hop-aligned index of buffer[0], so absolute → buffer-relative + // mel-frame is pcm_start_sample / hop. The first few buffer-relative + // frames are reflect-pad garbage but correspond to already-emitted + // absolute frames and are never read. const int64_t t_mel_start = ggml_time_us(); int mel_n_mels = 0; int mel_n_frames = 0; @@ -3620,9 +3058,8 @@ transcribe_status stream_feed( pc->n_threads); mst != TRANSCRIBE_OK) { - // For very short buffers (< 2 frames worth) compute returns - // INVALID_ARG; that's a "not enough audio yet" condition for - // streaming, not a fatal error. Treat as a no-op feed. + // For very short buffers compute returns INVALID_ARG — "not + // enough audio yet", a no-op feed, not a fatal error. if (mst == TRANSCRIBE_ERR_INVALID_ARG) { if (update != nullptr) { update->result_changed = false; @@ -3638,11 +3075,10 @@ transcribe_status stream_feed( } pc->t_mel_us += ggml_time_us() - t_mel_start; - // Emit as many chunks as we have new mel frames for. NeMo - // distinguishes first vs subsequent: first chunk has chunk_size_first - // NEW mel frames (no pre-encode-cache prepend, drop_extra=0); - // subsequent chunks have chunk_size_subsequent NEW frames plus a - // pre_encode_cache_size left-prepend (drop_extra=drop_extra_subsequent). + // Emit as many chunks as we have new mel frames for. First chunk: + // chunk_size_first new frames, no prepend, drop_extra=0. Subsequent: + // chunk_size_subsequent new + pre_encode_cache_size prepend, + // drop_extra=drop_extra_subsequent. const int pre_encode_cache_size = pm->hparams.enc_stream_pre_encode_cache_size > 0 ? pm->hparams.enc_stream_pre_encode_cache_size : pm->hparams.enc_subsampling_factor + 1; @@ -3670,21 +3106,11 @@ transcribe_status stream_feed( : pc->stream_caches.drop_extra_subsequent; const int prepend_frames = mel_fed - chunk_advance; - // Need (consumed + chunk_advance + right_edge_margin) total - // mel frames available before we can emit. The +margin is the - // right-edge reflect-pad zone of the mel transform: the last - // ceil(n_fft/2 / hop_length) frames computed at any moment - // have their windows reflect-padded against the buffer end, - // producing slightly wrong values until more audio arrives. - // Without this margin, chunk N's last few mel frames would - // get values from a "pre-mature" mel computation and never - // match NeMo's full-audio mel reference. The cost is a few - // mel frames of extra latency per chunk (40ms for nemotron's - // n_fft=512/hop=160). - // - // stream_finalize doesn't gate on this margin because by then - // we've received the EOF signal and the right-edge reflect-pad - // is "real" — it's the end of the stream. + // Hold back the last ceil(n_fft/2 / hop) mel frames: their STFT + // windows are reflect-padded against the buffer end and stay + // slightly wrong (never matching NeMo's full-audio mel) until more + // audio arrives. stream_finalize skips this margin since the + // right-edge pad is then real (end of stream). const int right_edge_margin = (pm->hparams.fe_n_fft / 2 + pm->hparams.fe_hop_length - 1) / pm->hparams.fe_hop_length; @@ -3705,10 +3131,7 @@ transcribe_status stream_feed( for (int m = 0; m < mel_n_mels; ++m) { // Pre-encode-cache prepend: [consumed - prepend_frames, - // consumed). On the first chunk prepend_frames == 0 so the - // loop is skipped entirely. On subsequent chunks consumed - // >= chunk_advance >= prepend_frames so the indices are - // never negative. + // consumed). prepend_frames == 0 on the first chunk. for (int h = 0; h < prepend_frames; ++h) { const int64_t frame_idx = consumed - prepend_frames + h; const float val = (frame_idx < 0) @@ -3724,11 +3147,10 @@ transcribe_status stream_feed( } } - // Flip is_first_chunk BEFORE the emit so the dump hooks see the - // correct step_num and drop semantics. The caller (emit) uses - // the parameters we already extracted. + // Flip is_first_chunk before the emit (the emit uses the params + // we already extracted). pc->stream_caches.is_first_chunk = false; - (void)pre_encode_cache_size; // kept for future first-chunk diagnostics + (void)pre_encode_cache_size; const int64_t t_enc_start = ggml_time_us(); if (const transcribe_status st = emit_streaming_chunk( @@ -3754,19 +3176,12 @@ transcribe_status stream_feed( pc->stream_revision += 1; } - // Trim already-consumed PCM. The next chunk's mel_at() reads back - // by `prepend_max` frames (the pre-encode-cache history) AND each - // mel-frame STFT window extends another `pad = n_fft/2` samples - // left of its center. Keep both margins or the prepend frames' - // windows reflect-pad against a truncated buffer start, silently - // corrupting the pre-encode-cache input on every subsequent chunk - // (observed as deletions at chunk boundaries in WER). - // Round the drop down to a hop boundary to keep pcm_start_sample - // aligned with mel-frame index 0. - // - // First-feed safety: with mel_frames_consumed == 0 the target - // keep_from is negative, the drop_aligned guard rejects it, and - // the buffer stays intact. + // Trim already-consumed PCM, keeping the prepend_max history frames + // AND each mel window's pad = n_fft/2 left margin — drop more and the + // prepend frames reflect-pad against a truncated start, corrupting the + // cache input (observed as chunk-boundary deletions in WER). Drop is + // hop-aligned so pcm_start_sample stays at mel-frame 0. First feed + // (mel_frames_consumed == 0): keep_from is negative, guard rejects. { const int pad = pm->hparams.fe_n_fft / 2; const int prepend_max = @@ -3798,12 +3213,9 @@ transcribe_status stream_feed( } } - // Audio "committed" cursor tracks the amount of audio we've fully - // consumed through the encoder (one mel frame = hop_length / - // sample_rate seconds). Using mel_frames_consumed (vs encoder - // output frame count) avoids the right-context overshoot that - // makes "encoder frames * 80 ms" exceed the actual audio time - // by a few ms per chunk on streaming variants. + // Audio "committed" cursor. mel_frames_consumed (vs encoder output + // frames) avoids the right-context overshoot that would make + // "encoder frames * 80 ms" exceed the real audio time per chunk. pc->stream_audio_committed_us = pc->stream_caches.mel_frames_consumed * static_cast(pm->hparams.fe_hop_length) * @@ -3838,12 +3250,9 @@ transcribe_status stream_finalize( // -------- Buffered streaming finalize -------- // - // Mirrors NeMo's reference loop tail exactly: one final emit that - // consumes all remaining audio (num_new = audio_len - next_audio_read) - // with is_last_chunk=true. add_frames_get_removed_ folds the existing - // samples_right slot plus this final num_new into the chunk slot, so - // the decoder gets every encoder frame past ctx_left. No zero-pad, - // no fixed-stride trailing chunks. + // One final emit consuming all remaining audio with is_last_chunk=true; + // add_frames_get_removed_ folds the right slot + this num_new into the + // chunk slot so the decoder gets every frame past ctx_left (no zero-pad). if (pc->buf_active) { const int64_t total = static_cast(pc->stream_pcm_buffer.size()); @@ -3879,12 +3288,9 @@ transcribe_status stream_finalize( // -------- Cache-aware streaming finalize -------- // - // Recompute mel one last time and check for a sub-chunk tail. - // If there are any unprocessed mel frames left over (i.e. the - // stream ended mid-chunk), zero-pad to chunk_size_total and run a - // final chunk so we don't drop the tail audio. Sub-chunk tails - // produce a smaller post-subsample/drop output, but the encoder - // graph handles arbitrary mel chunk sizes. + // Recompute mel one last time; if the stream ended mid-chunk, run a + // final partial chunk so the tail audio isn't dropped (the encoder + // graph handles arbitrary mel chunk sizes). if (!pc->stream_pcm_buffer.empty()) { int mel_n_mels = 0; int mel_n_frames = 0; @@ -3917,14 +3323,9 @@ transcribe_status stream_finalize( ? pc->stream_caches.chunk_size_first : pc->stream_caches.chunk_size_subsequent; const int prepend_frames = mel_fed_full - chunk_advance_full; - // Feed the natural partial-chunk size — no silence pad. - // Matches NeMo's last-chunk behavior (the streaming buffer - // gives whatever's left; conformer_stream_step happily - // accepts a shorter chunk). The encoder graph rebuilds - // each call so it's fine to vary the shape. The decoder - // sees fewer enc_out frames but the trailing-audio - // tokens (e.g. the closing punctuation) survive instead - // of getting masked by zero-pad frames. + // Feed the natural partial-chunk size, no silence pad + // (NeMo's last-chunk behavior), so trailing-audio tokens + // survive instead of being masked by zero-pad frames. const int new_take = static_cast(remaining); const int mel_fed = prepend_frames + new_take; std::vector chunk( @@ -3991,16 +3392,12 @@ transcribe_status stream_finalize( void stream_reset(transcribe_session * session) { auto * pc = static_cast(session); - // Drop the buffered audio contents but keep the allocation for - // the next stream. - pc->stream_pcm_buffer.clear(); + pc->stream_pcm_buffer.clear(); // keep the allocation } -// Kind+slot probe. Parakeet ships no run-slot extensions, so the _RUN -// slot is always false. On the _STREAM slot, cache-aware variants -// (ChunkedLimited) take the PARAKEET_STREAM kind; chunked-attention -// variants (ChunkedLimitedWithRc) take the PARAKEET_BUFFERED_STREAM -// kind. Offline-only Parakeet variants accept neither. +// Kind+slot probe. No run-slot extensions (always false on _RUN). On +// _STREAM: ChunkedLimited takes PARAKEET_STREAM, ChunkedLimitedWithRc +// takes PARAKEET_BUFFERED_STREAM, offline variants neither. bool accepts_ext_kind(const transcribe_model * model, transcribe_ext_slot slot, uint32_t kind) { @@ -4020,9 +3417,8 @@ bool accepts_ext_kind(const transcribe_model * model, } // namespace -// `extern const` to force external linkage. Without `extern`, a -// namespace-scope `const` object has internal linkage in C++ and the -// forward declaration in transcribe-arch.cpp would not resolve. +// `extern const` forces external linkage (a namespace-scope const object +// is internal-linkage in C++ otherwise). extern const Arch arch = { /* .name = */ "parakeet", /* .load = */ load, @@ -4041,10 +3437,8 @@ extern const Arch arch = { // --------------------------------------------------------------------------- // Public parakeet extension init functions (global scope, C linkage). -// Defined here in family source so transcribe.cpp stays family-agnostic. -// Each stamps the transcribe_ext header (size + kind) and the field -// defaults; they are the single source of truth now that the INIT macros -// are gone. +// Defined here so transcribe.cpp stays family-agnostic; each stamps the +// transcribe_ext header (size + kind) and the field defaults. // --------------------------------------------------------------------------- extern "C" void transcribe_parakeet_stream_ext_init( diff --git a/src/arch/parakeet/parakeet.h b/src/arch/parakeet/parakeet.h index 64f987b4..b8a94afe 100644 --- a/src/arch/parakeet/parakeet.h +++ b/src/arch/parakeet/parakeet.h @@ -1,16 +1,9 @@ // arch/parakeet/parakeet.h - Parakeet-family internal model and context -// types. -// -// This header is INTERNAL to the Parakeet sources. It defines the -// concrete classes that derive from transcribe_model / transcribe_session -// for this family. It is also visible to tests/transcribe_tokenizer_smoke -// (via the test target's PRIVATE include path) so the test can pull a -// const Tokenizer * out of a model loaded through the public C ABI -// without inventing a temporary public accessor. -// -// Other source files outside src/arch/parakeet/ should not include this -// header. The central dispatcher in transcribe.cpp talks to the family -// only through the Arch trait and the base classes. +// types: the concrete classes deriving from transcribe_model / +// transcribe_session. Internal to the Parakeet sources (also visible to +// tests/transcribe_tokenizer_smoke via the test target's PRIVATE include +// path). The central dispatcher talks to the family only through the Arch +// trait and the base classes. #pragma once @@ -38,42 +31,25 @@ typedef struct ggml_backend_sched * ggml_backend_sched_t; namespace transcribe::parakeet { -// --------------------------------------------------------------------------- -// Attention masks -// --------------------------------------------------------------------------- -// -// Host-side helpers that materialize the 2D attention masks used by the -// streaming paths. Declared here (in the family-internal header) so the -// per-mask unit tests can call them directly without needing a ggml -// backend or a real GGUF on disk. -// -// chunked_limited_with_rc mask (buffered streaming, parakeet-unified-en-0.6b) -// -// Mirrors NeMo `ConformerEncoder._create_masks` lines 843-869 for the -// chunked_limited_with_rc branch: -// -// c_q = q / C -// window_start = max(0, c_q * C - L) -// window_end = min(T - 1, c_q * C + C - 1 + R) -// mask[q, k] = 0 if window_start <= k <= window_end else -INF +// Host-side attention-mask helpers for the streaming paths. Declared +// here so the per-mask unit tests can call them without a ggml backend +// or a GGUF on disk. // -// The output buffer is row-major [T_q, T_k] with T_q == T_k == T — -// the buffered driver always runs the encoder over a power-of-T -// window, so query and key axes match. out_buf must point to at -// least T*T floats. The mask broadcasts over heads in rel_pos_mhsa. +// chunked_limited_with_rc mask (buffered streaming). Mirrors NeMo +// ConformerEncoder._create_masks: +// c_q = q / C +// window_start = max(0, c_q * C - L) +// window_end = min(T - 1, c_q * C + C - 1 + R) +// mask[q, k] = 0 if window_start <= k <= window_end else -INF +// Output is row-major [T_q, T_k] with T_q == T_k == T (>= T*T floats); +// broadcasts over heads in rel_pos_mhsa. // -// Asserts (debug-only): T >= 1, C >= 1, L >= 0, R >= 0. -// -// `pad_length` (optional) folds in NeMo's conv-overhang pad_mask: -// cells where row q >= pad_length OR col k >= pad_length are masked -// regardless of the chunked window. This matches NeMo's -// ConformerEncoder.forward, which ANDs pad_mask onto the chunked -// mask before MHA. Set pad_length = T to disable padding masking -// (the offline / no-pad-overhang case). For buffered streaming, -// pass effective_T = ctx.total / encoder_frame2audio_samples; the -// conv frontend may produce T_enc > effective_T due to padding -// overhang, and those trailing frames must be masked out or they -// contaminate the attention scores of valid frames. +// `pad_length` (optional) folds in NeMo's conv-overhang pad_mask: cells +// where q >= pad_length OR k >= pad_length are masked regardless of the +// window. pad_length = T disables it (offline). For buffered streaming +// pass effective_T = ctx.total / encoder_frame2audio_samples — the conv +// frontend may produce T_enc > effective_T, and those trailing frames +// must be masked or they contaminate valid frames' attention scores. void compute_chunked_limited_with_rc_mask( float * out_buf, int T, @@ -83,42 +59,26 @@ void compute_chunked_limited_with_rc_mask( int pad_length); -// Family defaults — applied before transcribe::read_capability_kv runs, -// so capability KV present overrides the default and capability KV -// absent leaves the default in place. Defined in capabilities.cpp. -// -// There is no per-variant resolver: variant identity is descriptive -// metadata, not a behavioral discriminator. Differences between v2 -// and v3 (or any future Parakeet variant) are expressed as -// stt.capability.* / general.languages KV that the converter writes -// and the family-agnostic readers in transcribe-meta.h consume. -// RESUME.md "Decisions still load-bearing" has the rationale and the -// llama.cpp / whisper.cpp comparison that drove it. +// Family defaults — applied before transcribe::read_capability_kv runs +// (KV present overrides, KV absent leaves the default). Defined in +// capabilities.cpp. There is no per-variant resolver: variant identity +// is descriptive metadata; v2/v3 differences are expressed as +// stt.capability.* / general.languages KV. void apply_family_invariants(transcribe_model & model); -// Concrete model. Owns the ggml_context that holds every weight -// tensor's data buffer (allocated by gguf_init_from_file with -// no_alloc=false during load). The destructor frees the ggml_context; -// every borrowed ggml_tensor* in `weights` is invalidated at that -// moment. The first gguf_context (the loader's header-only inspection -// pass) is no longer kept alive past load() — its only consumer is -// the tokenizer ingest, which copies all the data it needs. The -// second gguf_context (the one created by load() to read tensor -// metadata) is freed inside load() once the weights have been -// catalogued. +// Concrete model. Owns ctx_meta, which holds every weight tensor's +// data buffer; the destructor frees it, invalidating every borrowed +// ggml_tensor* in `weights`. struct ParakeetModel final : public transcribe_model { Tokenizer tok; ParakeetHParams hparams; ParakeetWeights weights; ggml_context * ctx_meta = nullptr; - // Runtime backend plan. Resolved at load() time from - // transcribe_model_load_params::backend via - // transcribe::load_common::init_backends. See - // transcribe-backend.h for the plan semantics. The plan's - // `scheduler_list` owns the backends; the destructor frees - // them in reverse order. Helpers that key off "is this CPU?" - // check `plan.primary_kind == BackendKind::Cpu` directly. + // Runtime backend plan, resolved at load() via + // load_common::init_backends (see transcribe-backend.h). plan's + // scheduler_list owns the backends; the destructor frees them in + // reverse order. transcribe::BackendPlan plan; ggml_backend_buffer_t backend_buffer = nullptr; @@ -127,48 +87,26 @@ struct ParakeetModel final : public transcribe_model { ggml_context * bn_fused_ctx = nullptr; ggml_backend_buffer_t bn_fused_buffer = nullptr; - // On a CPU primary backend the conformer 1×1 pointwise conv weights - // are dequantized from their on-disk F16 form to F32 at load time. - // The current quantizer (transcribe-quantize, post commit 6fee9b9) - // routes every Conformer ConvPw to F16 across all presets so that - // GPU backends with native F16 compute (Metal, Vulkan, CUDA) get - // the bandwidth halving for free; on CPUs without native F16 - // arithmetic (Zen 2 and earlier) the per-matmul F16→F32 upconvert - // erases the bandwidth win and dwarfs the original cost. Promoting - // back to F32 at load time pays the conversion exactly once. - // - // Today's parakeet GGUFs predate the universal F16 conv_pw policy - // and ship F32 conv_pw, so the promotion is a no-op for them. The - // wiring is in place so the next regen does the right thing - // automatically. The promotion is also a no-op on every non-CPU - // primary backend. - // - // The CohereBlock comment about ~235 MB of "wasted" originals - // applies symmetrically: ggml backend buffers don't support - // freeing individual tensor slots, so the F16 source tensors stay - // resident in the main weight buffer after promotion. The wired - // weight slot is repointed at the F32 copy and the encoder graph - // never touches the originals again. + // On a CPU primary backend, the conformer 1×1 pointwise conv weights + // are dequantized from F16 to F32 at load: CPUs without native F16 + // arithmetic (Zen 2 and earlier) pay a per-matmul F16→F32 upconvert + // that erases the bandwidth win, so promoting once at load is cheaper. + // No-op on GPU backends and on non-CPU primary backends. The F16 + // source tensors stay resident (ggml can't free individual slots); + // the weight slot is repointed at the F32 copy. ggml_context * conv_pw_f32_ctx = nullptr; ggml_backend_buffer_t conv_pw_f32_buffer = nullptr; - // Mel front-end. Constructed once at load() time from the - // hparams (precomputes the periodic Hann window + Slaney mel - // filterbank). MelFrontend is documented as const-after-ctor - // and thread-safe across compute() calls, so it's shared by - // every context derived from this model. std::optional because - // MelFrontend has no default constructor; emplace() in load() - // after the hparams are read. + // Mel front-end, constructed once at load() from the hparams + // (precomputes Hann window + Slaney mel filterbank). + // const-after-ctor and thread-safe, so shared by every context. + // std::optional because MelFrontend has no default constructor. std::optional mel; - // Host mirror of the predictor + joint weights, populated at - // load() time from the backend tensors via - // build_host_decoder_weights. The decoder runs on host (see - // decoder.h for the rationale); having a const-after-load - // mirror lets every context share the same arrays without - // each transcribe_run paying a backend readback cost. The - // memory cost is small relative to the encoder weights - // (~35 MB on v2, ~73 MB on v3 vs ~2.4 GB total). + // Host mirror of the predictor + joint weights, populated at load() + // via build_host_decoder_weights. The decoder runs on host (see + // decoder.h); the const-after-load mirror lets every context share it + // without a per-run backend readback. HostDecoderWeights host_decoder; ParakeetModel() = default; @@ -177,32 +115,17 @@ struct ParakeetModel final : public transcribe_model { const transcribe::Tokenizer * tokenizer() const override { return &tok; } }; -// Per-context streaming encoder state. Allocated lazily on the first +// Per-context streaming encoder state, allocated lazily on the first // stream_begin for ChunkedLimited variants. Mirrors NeMo's -// cache_last_channel + cache_last_time + cache_last_channel_len tuple -// returned by get_initial_cache_state. Layout cheat sheet: -// -// cache_last_channel [d_model, T_cache, n_layer] f32 -// Per-layer post-attn-LN input cache. Concat-prepended to x_norm -// before Q/K/V projection in the streaming MHSA path. NeMo stores -// the pre-projection tensor (not split K/V), so a single per-layer -// slot suffices. -// -// cache_last_time [k_minus_1, d_model, n_layer] f32 -// Per-layer post-pw1+GLU conv-module input cache. Replaces the -// zero-pad on the left of the depthwise conv in the streaming -// conv_module path. -// -// cache_last_channel_len scalar -// Valid-frame count in cache_last_channel (0..T_cache). Drives -// the streaming-mask offset so unfilled prefix rows are masked -// out until the cache fills. -// -// Sizes for nemotron-speech-streaming-en-0.6b at att_context_size=[70,13]: -// T_cache = 70 (att_context_left) -// k_minus_1 = 8 (conv_kernel - 1) -// n_layer = 24, d_model = 1024 -// Total cache buffer ~7.5 MB. +// get_initial_cache_state tuple. Layout cheat sheet: +// cache_last_channel [d_model, T_cache, n_layer] f32 — per-layer +// post-attn-LN cache, concat-prepended to x_norm before Q/K/V. +// cache_last_time [k_minus_1, d_model, n_layer] f32 — per-layer +// post-pw1+GLU cache, replaces the depthwise conv's left zero-pad. +// cache_last_channel_len scalar — valid frames (0..T_cache); drives +// the streaming-mask offset until the cache fills. +// Sizes for nemotron-speech-streaming-en-0.6b at [70,13]: T_cache=70, +// k_minus_1=8, n_layer=24, d_model=1024 (~7.5 MB total). struct ParakeetStreamingCaches { ggml_context * ctx = nullptr; ggml_backend_buffer_t buffer = nullptr; @@ -213,27 +136,21 @@ struct ParakeetStreamingCaches { std::vector last_channel; std::vector last_time; - // KV-cache variant of last_channel: per-layer pre-projected - // attention keys/values (bias included) over the same T_cache - // window, ne[d_model, T_cache, 1, 1] each. Mathematically - // equivalent to recomputing K/V from last_channel every chunk — - // k(t)/v(t) depend only on the frame's pre-attention activation — - // but the per-chunk K/V projections then run on T_q_new columns - // instead of T_cache + T_q_new. The default streaming graph - // consumes these and leaves last_channel untouched; an active dump - // harness (which compares last_channel against NeMo) falls back to - // the channel recompute path. Cleared alongside the other caches at + // KV-cache variant of last_channel: per-layer pre-projected K/V + // (bias included) over the T_cache window, ne[d_model, T_cache, 1, 1] + // each. Equivalent to recomputing K/V from last_channel every chunk + // but the per-chunk projections run on T_q_new columns instead of + // T_cache + T_q_new. The default graph consumes these; an active dump + // harness falls back to the channel recompute path. Cleared at // stream_begin. std::vector last_k; std::vector last_v; // Rel-pos projection memoization: per-layer // [head_dim, pos_proj_len, n_head, 1] f32 tensors holding - // attn_pos_w @ pos_emb in attention-ready layout. Geometry-keyed - // (pos_proj_len), not content-keyed, so it survives across chunks - // AND streams; rebuilt by ensure_pos_proj_cache whenever a chunk's - // pos_len differs. Own ctx + buffer because the size changes with - // geometry while the caches above are fixed at stream_begin. + // attn_pos_w @ pos_emb. Geometry-keyed (pos_proj_len), so it survives + // across chunks and streams; rebuilt by ensure_pos_proj_cache on a + // pos_len change. Own ctx + buffer (size varies with geometry). ggml_context * pos_proj_ctx = nullptr; ggml_backend_buffer_t pos_proj_buf = nullptr; std::vector pos_proj; @@ -241,45 +158,26 @@ struct ParakeetStreamingCaches { int32_t channel_len = 0; - // Absolute mel-frame cursor across the whole stream — used to - // anchor token / segment timestamps to original-audio time even - // when chunks vary in size. Counts mel frames consumed (i.e. fed - // to the encoder, minus history prepend). + // Absolute mel-frame cursor across the whole stream, anchoring + // token/segment timestamps to original-audio time even when chunks + // vary in size. Counts mel frames consumed (minus history prepend). int64_t mel_frames_consumed = 0; // Absolute sample index of stream_pcm_buffer[0]. The buffer is a - // sliding window: after each emit, prefix samples no longer needed - // by any unemitted frame's STFT window are trimmed (hop-aligned). - // Bounded per-feed mel cost: O(new audio) instead of O(total). - // Always a multiple of hop_length so mel-frame indexing translates - // cleanly between absolute and buffer-relative. + // sliding window: prefix samples no longer needed by any unemitted + // frame's STFT window are trimmed (hop-aligned), bounding per-feed mel + // cost to O(new audio). Always a hop_length multiple so mel-frame + // indexing translates cleanly between absolute and buffer-relative. int64_t pcm_start_sample = 0; - // Streaming geometry for the active stream. Resolved at - // stream_begin from ParakeetHParams + the caller-selected - // att_context_right (transcribe_parakeet_stream_ext). All four - // fields are constant across chunks within one stream; they - // change only when stream_begin is called again with a different - // att_context_right selection. - // - // att_context_right: chosen R from the model's training menu - // (default = enc_att_context_right). - // att_context_left: matching L (currently always the model - // default; multi-L selection isn't in the API). - // - // First-chunk and subsequent-chunk geometry differ. NeMo's - // setup_streaming_params packs the (first, subsequent) pair into - // `streaming_cfg.chunk_size` and `shift_size`; we mirror that - // with parallel fields below. The driver in stream_feed picks - // the right pair based on the (now-tracked) is_first_chunk flag. - // - // For nemotron-speech-streaming-en-0.6b at att_context_right=13: - // chunk_size_first = 105 (sampling_frames_first + 8*R) - // chunk_size_subsequent = 112 (subsampling_factor * (1 + R)) - // mel_fed_first = 105 (no pre_encode_cache prepend) - // mel_fed_subsequent = 121 (9 cache + 112 new) - // drop_extra_first = 0 - // drop_extra_subsequent = 2 + // Streaming geometry, resolved at stream_begin from ParakeetHParams + + // the caller-selected att_context_right. Constant across chunks within + // a stream. First- vs subsequent-chunk geometry differ (NeMo's + // setup_streaming_params); the driver picks the pair via + // is_first_chunk. For nemotron-speech-streaming-en at R=13: + // chunk_size_first=105, chunk_size_subsequent=112, + // mel_fed_first=105, mel_fed_subsequent=121 (9 cache + 112 new), + // drop_extra_first=0, drop_extra_subsequent=2. int att_context_right = 0; int att_context_left = 0; int chunk_size_first = 0; @@ -288,141 +186,88 @@ struct ParakeetStreamingCaches { int mel_fed_subsequent = 0; int drop_extra_first = 0; int drop_extra_subsequent = 0; - // Whether the next emit_streaming_chunk call is the FIRST in this - // stream. Flips to false after the first chunk is processed. + // Whether the next emit_streaming_chunk is the FIRST in this stream. bool is_first_chunk = true; - // Per-stream chunk counter (== "step_num" in NeMo's - // perform_streaming reference). Resets to 0 on stream_begin, then - // increments once per emit_streaming_chunk. Used for indexing dump - // filenames against the Python reference dumps when - // TRANSCRIBE_DUMP_DIR is set. + // Per-stream chunk counter (NeMo's step_num). Resets at stream_begin, + // increments per emit; indexes dump filenames. int chunk_step = 0; bool initialized = false; }; -// Per-context streaming RNN-T decoder state. The offline path -// allocates LstmState locally; for streaming we own the state on the -// context so it carries across stream_feed calls. Reset on each -// stream_begin to a fresh "start of sequence" state. +// Per-context streaming RNN-T decoder state, owned on the context so it +// carries across stream_feed calls. Reset at each stream_begin to a fresh +// start-of-sequence state. struct ParakeetStreamingDecoderState { - // LSTM (h, c) per layer + scratch buffers. Sized once on - // stream_begin from host_decoder.predictor.n_layers / - // pred_hidden. + // LSTM (h, c) per layer; sized at stream_begin. LstmState lstm_state; - // Last emitted non-blank token id; the predictor input. Starts at - // a sentinel (-1) meaning "predictor seeded with SOS embedding". + // Last emitted non-blank token id (predictor input); -1 = SOS. int32_t prev_token_id = -1; - // Absolute encoder-frame offset for converting per-chunk - // step_at_emit indices into stream-wide frame indices. + // Absolute encoder-frame offset for converting per-chunk step_at_emit + // into stream-wide frame indices. int64_t frame_offset = 0; bool initialized = false; }; // Concrete context. Owns a per-call compute context and a persistent -// multi-backend scheduler that dispatches encoder graph ops to the -// best available backend (GPU for matmuls, BLAS for CPU matmuls, -// CPU for everything else). +// multi-backend scheduler that dispatches encoder graph ops to the best +// available backend. struct ParakeetSession final : public transcribe_session { - // Compute context: holds the cgraph and intermediate tensor - // metadata. Created with no_alloc=true; data lives in the - // sched-managed buffers. Reset at the start of every run() call. + // Compute context: cgraph + intermediate tensor metadata. no_alloc; + // data lives in sched-managed buffers. Reset each run(). ggml_context * compute_ctx = nullptr; // Multi-backend scheduler. Persists across calls; manages compute - // buffer allocation and dispatches ops to the best backend. - // Reuses buffers when the graph topology is unchanged. + // buffer allocation and reuses buffers when topology is unchanged. ggml_backend_sched_t sched = nullptr; - // Encoder forward output, borrowed pointer into compute_ctx - // tensors. Invalidated when compute_ctx is reset on the next - // run() call. + // Encoder forward output, borrowed into compute_ctx; invalidated when + // compute_ctx is reset next run(). ggml_tensor * encoder_out = nullptr; - // Per-context scratch reused across runs. Keeping these on the - // context avoids re-allocating the same medium-sized host buffers - // every call while preserving the exact arithmetic. + // Per-context scratch reused across runs. std::vector mel_buf; std::vector pos_buf; std::vector pos_div_term; std::vector enc_host; std::vector raw_tokens; - // Per-call timings (t_mel_us, t_encode_us, t_decode_us) live - // on the base transcribe_session now and are surfaced via the - // public transcribe_get_timings API. - - // Phase 5: TDT decode result storage lives on the base - // (transcribe_session::tokens / words / segments / full_text / - // result_kind / has_result), populated by Parakeet::run during - // decode. No per-family result fields here. + // Per-call timings and the TDT decode result (tokens / words / + // segments / full_text / result_kind / has_result) live on the base + // transcribe_session; no per-family copies here. - // KV type for flash attention, resolved from the context params. - - // ---- streaming-of-whole state (Phase 4a) ---- - // - // Parakeet exposes the streaming API surface for cache-aware - // streaming variants (today: nemotron-speech-streaming-en-0.6b, - // which has enc_att_context_style=ChunkedLimited). Phase 4a is a - // stream-of-whole stub: feed appends PCM, finalize runs the - // existing one-shot inference helper on the accumulated buffer. - // This makes the dispatcher / lifecycle / result-snapshot path - // observable end-to-end and trivially parity-equivalent to - // transcribe_run on the same audio; real per-chunk encoder - // feeding (cache_last_channel / cache_last_time / RNNT predictor - // carry) is Phase 4b. - // - // Lifecycle: - // stream_begin: stream_pcm_buffer.clear(); stream_run_params = *run_params; - // stream_feed: append samples - // stream_finalize: drain buffer through run_one_shot_inner() - // stream_reset: stream_pcm_buffer.clear() (keep capacity) + // ---- streaming-of-whole state ---- // - // These fields are NOT touched by clear_result; the dispatcher - // wipes lifecycle-agnostic snapshot state there, and the family - // owns its per-utterance audio scratch. + // stream_begin captures run_params; stream_feed appends PCM; + // stream_finalize drains the buffer through the inference helper; + // stream_reset clears (keeping capacity). NOT touched by + // clear_result (the family owns its per-utterance audio scratch). std::vector stream_pcm_buffer; transcribe_run_params stream_run_params {}; - // ---- M2 incremental streaming state (cache-aware) ---- + // ---- incremental streaming state (cache-aware) ---- // - // Allocated on the first stream_begin for ChunkedLimited variants - // (today: nemotron-speech-streaming-en-0.6b). Zeroed at each - // stream_begin so each utterance starts from a clean cache. - // Persists across feed/finalize within an utterance; freed in the - // context destructor. + // Allocated on the first stream_begin for ChunkedLimited variants, + // zeroed at each stream_begin, persists across feed/finalize within + // an utterance, freed in the dtor. ParakeetStreamingCaches stream_caches; ParakeetStreamingDecoderState stream_dec_state; // ---- Buffered streaming state (parakeet-unified-en-0.6b) ---- // - // Mirrors NeMo's `StreamingBatchedAudioBuffer` + the reference - // inference loop at - // `examples/asr/asr_chunked_inference/rnnt/speech_to_text_streaming_infer_rnnt.py`. - // Variable-stride: step 0 (initial fill) consumes - // `samples_chunk + samples_right` of new audio; subsequent feed - // steps consume `samples_chunk`; finalize's single last step - // consumes whatever's left (chunk slot absorbs the trailing right - // context, no zero-pad). - // - // The expected geometry (buf_left_frames / buf_samples_left etc.) - // is constant per stream and reflects the active (L, C, R) tuple. - // The ctx_* fields track the buffer's INTERNAL `ContextSize` - // (left/chunk/right slot sizes in samples), updated per-chunk via - // `buf_ctx_add_frames` to mirror NeMo's `add_frames_get_removed_`. - // The encoder output sliced off this chunk = (ctx_left / - // frame_samples) leading frames; the decoder decodes - // (ctx_chunk / frame_samples) frames on non-last, or all remaining - // frames on last. - // - // The RNN-T LSTM state + last-emitted token id ride on the - // existing stream_dec_state slot — buffered streaming reuses the - // same predictor/joint code path the cache-aware driver does, - // just without a per-layer encoder cache. + // Mirrors NeMo's StreamingBatchedAudioBuffer + reference inference + // loop. Variable-stride: step 0 consumes samples_chunk + samples_right + // of new audio; subsequent feeds consume samples_chunk; finalize's + // last step consumes the rest (chunk slot absorbs the trailing right + // context, no zero-pad). The buf_* geometry reflects the active + // (L, C, R) tuple; the ctx_* fields track the buffer's internal + // ContextSize, updated per-chunk via buf_ctx_add_frames (NeMo's + // add_frames_get_removed_). The RNN-T state rides on stream_dec_state + // (same predictor/joint path, no per-layer encoder cache). int32_t buf_left_frames = 0; // L (expected) int32_t buf_chunk_frames = 0; // C int32_t buf_right_frames = 0; // R diff --git a/src/arch/parakeet/weights.cpp b/src/arch/parakeet/weights.cpp index c1803e53..7035d32b 100644 --- a/src/arch/parakeet/weights.cpp +++ b/src/arch/parakeet/weights.cpp @@ -1,18 +1,10 @@ -// arch/parakeet/weights.cpp - implementation of read_parakeet_hparams -// and build_parakeet_weights. +// arch/parakeet/weights.cpp - read_parakeet_hparams and +// build_parakeet_weights. // -// Pattern follows llama.cpp's per-arch load_tensors function: -// - read every required hparam from KV explicitly -// - then build the tensor catalog as a sequence of get_tensor() -// calls with explicit expected shapes -// - missing tensor or shape mismatch -> TRANSCRIBE_ERR_GGUF with a -// log naming the offending tensor -// -// The shape contract documented in weights.h is enforced here. The -// converter has to write tensors in these shapes; the synthetic -// fixture in tests/fixtures/make_gguf_fixtures.py emits them in -// these shapes; phase 4's encoder builder will read them in these -// shapes. +// Read every required hparam from KV explicitly, then build the tensor +// catalog as get_tensor() calls with explicit expected shapes; a missing +// tensor or shape mismatch returns TRANSCRIBE_ERR_GGUF naming the tensor. +// The shape contract documented in weights.h is enforced here. #include "weights.h" @@ -57,31 +49,21 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, if (auto st = read_required_u32_kv(gguf, "stt.parakeet.encoder.subsampling_channels", kFamilyTag, hp.enc_subsampling_channels); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.parakeet.encoder.pos_emb_max_len", kFamilyTag, hp.enc_pos_emb_max_len); st != TRANSCRIBE_OK) return st; - // Optional. NeMo Parakeet v2/v3 ship with use_bias=false; the - // 1.1B and rnnt/ctc variants ship with use_bias=true (11 biases per - // block). Default matches v2/v3 so legacy GGUFs without the KV stay - // on the bias-free path. + // use_bias: false for v2/v3, true for 1.1B and rnnt/ctc (11 biases + // per block). Default false for legacy GGUFs without the KV. if (auto st = read_optional_bool_kv(gguf, "stt.parakeet.encoder.use_bias", kFamilyTag, false, hp.enc_use_bias); st != TRANSCRIBE_OK) return st; - // xscaling: NeMo's RelPositionalEncoding multiplies x by - // sqrt(d_model) before generating the pos_emb when this is true. - // v2/v3/tdt-* are False; ctc-*/rnnt-*/unified-en are True. - // Default to false so legacy GGUFs (which don't ship this KV) - // continue to work without scaling - they're all xscaling=False - // models in the current set. + // xscaling (see weights.h). Default false for legacy GGUFs. if (auto st = read_optional_bool_kv(gguf, "stt.parakeet.encoder.xscaling", kFamilyTag, false, hp.enc_xscaling); st != TRANSCRIBE_OK) return st; - // Local attention window. Default -1/-1 = full attention; matches - // every variant except parakeet-tdt_ctc-1.1b ([128, 128] regular - // local) and nemotron-speech-streaming-en-0.6b ([70, 13] chunked). + // Local attention window. Default -1/-1 = full attention. if (auto st = read_optional_int32_kv(gguf, "stt.parakeet.encoder.att_context_left", kFamilyTag, -1, hp.enc_att_context_left); st != TRANSCRIBE_OK) return st; if (auto st = read_optional_int32_kv(gguf, "stt.parakeet.encoder.att_context_right", kFamilyTag, -1, hp.enc_att_context_right); st != TRANSCRIBE_OK) return st; - // Multi-lookahead training menu (cache-aware streaming models). Flat - // int32 array [L0,R0,L1,R1,...]; index 0 must equal (att_context_left, - // att_context_right). Optional — absent for offline GGUFs, in which - // case we synthesize a one-element list from the scalar fields so - // call sites read uniformly. + // Multi-lookahead training menu (cache-aware streaming). Flat int32 + // [L0,R0,L1,R1,...]; index 0 must equal (att_context_left, + // att_context_right). Absent for offline GGUFs → synthesize a + // one-element list from the scalar fields. { std::vector flat; switch (read_int32_array_kv(gguf, "stt.parakeet.encoder.att_context_size_choices", flat)) { @@ -126,11 +108,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } - // Chunked-limited-with-rc training menu (buffered streaming variants). - // Three independent int32 arrays carrying the L / C / R choices the - // model was trained against. Absent for non-buffered-streaming GGUFs; - // when present they're a structured triple read into the hparams - // and the runtime picks one entry from each at stream_begin time. + // Chunked-limited-with-rc training menu (buffered streaming). Three + // independent int32 L / C / R arrays; absent for non-buffered GGUFs. { auto read_menu = [&](const char * key, std::vector & out) -> transcribe_status { std::vector raw; @@ -160,9 +139,7 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, st != TRANSCRIBE_OK) return st; } - // Attention context style. Optional with default "regular" so legacy - // GGUFs (every variant before nemotron-speech-streaming-en-0.6b) - // stay on the existing local/full path. + // Attention context style. Optional, default "regular" for legacy GGUFs. { std::string style; if (auto st = read_optional_string_kv(gguf, "stt.parakeet.encoder.att_context_style", @@ -199,10 +176,9 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } - // Conv module: per-side depthwise padding and norm type. Both are - // optional with defaults that match the offline-Conformer config — - // -1 means "centred (k-1)/2"; "batch_norm" means the historic - // fused BN path with running_mean / running_var. + // Conv module: per-side depthwise padding and norm type. Optional; + // defaults -1 (centred (k-1)/2) and "batch_norm" (fused BN with + // running_mean / running_var). if (auto st = read_optional_int32_kv(gguf, "stt.parakeet.encoder.conv_context_left", kFamilyTag, -1, hp.enc_conv_context_left); st != TRANSCRIBE_OK) return st; if (auto st = read_optional_int32_kv(gguf, "stt.parakeet.encoder.conv_context_right", kFamilyTag, -1, hp.enc_conv_context_right); st != TRANSCRIBE_OK) return st; { @@ -226,9 +202,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } - // Cache-aware streaming pre-encode constants. Optional — both KVs - // are emitted only by streaming-trained variants. Default 0 (means - // "non-streaming"); the streaming code paths gate on both > 0. + // Cache-aware streaming pre-encode constants. Optional, default 0 + // (non-streaming); the streaming paths gate on both > 0. if (auto st = read_optional_int32_kv(gguf, "stt.parakeet.encoder.streaming.pre_encode_cache_size", kFamilyTag, 0, hp.enc_stream_pre_encode_cache_size); st != TRANSCRIBE_OK) @@ -248,9 +223,7 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, return st; } - // Head kind dispatch. Optional KV with default "tdt" so legacy - // v2/v3 GGUFs (predate the KV) resolve to TDT, which is what they - // are. The 8 new variants all ship the KV explicitly. + // Head kind dispatch. Optional KV, default "tdt" for legacy v2/v3. { std::string head_kind_str; if (auto st = read_optional_string_kv(gguf, "stt.parakeet.head_kind", @@ -287,11 +260,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } if (hp.head_kind == HeadKind::TDT) { - // TDT decoding parameters. `durations` is required (the decoder - // can't run without it); `max_symbols` is optional with a default - // matching every published v2/v3 config we've seen. RNNT does - // NOT carry these — its joint emits `vocab+1` (no extras) and - // the decode loop has no duration choice. + // TDT decoding parameters. durations required; max_symbols + // optional (default 10). RNNT carries neither. switch (read_int32_array_kv(gguf, "stt.parakeet.tdt.durations", hp.tdt_durations)) { case KvResult::Ok: break; @@ -317,10 +287,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } } else if (hp.head_kind == HeadKind::RNNT) { - // RNNT shares the predictor + joint code paths with TDT, but has - // no durations and the joint emits exactly `vocab+1` logits. - // Reuse `tdt_max_symbols` as the decode loop's per-frame stuck - // cap (NeMo's RNNTGreedyDecodeBatched uses the same constant). + // RNNT: no durations, joint emits exactly vocab+1. tdt_max_symbols + // is reused as the per-frame stuck cap. hp.tdt_durations.clear(); hp.tdt_max_symbols = 10; } else { @@ -329,15 +297,10 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, hp.tdt_max_symbols = 0; } - // Frontend. The full stt.frontend.* block PLAN.md declares as the - // complete contract between converter and loader. Reading every - // field at load time, even though phase 3's C++ frontend hasn't - // landed yet, makes the loader the gate that catches a future - // converter that drops a field. Required keys: type, num_mels, - // sample_rate, n_fft, win_length, hop_length, window, normalize, - // dither, pre_emphasis, f_min, f_max. CMVN/LFR fields are - // intentionally not read here — Parakeet doesn't use them and - // the converter doesn't emit them. + // Frontend. The complete stt.frontend.* contract between converter + // and loader; reading every field makes the loader the gate that + // catches a converter that drops one. CMVN/LFR fields are not read + // (Parakeet doesn't use them). if (auto st = read_required_string_kv(gguf, "stt.frontend.type", kFamilyTag, hp.fe_type); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.num_mels", kFamilyTag, hp.fe_num_mels); st != TRANSCRIBE_OK) return st; if (auto st = read_required_u32_kv(gguf, "stt.frontend.sample_rate", kFamilyTag, hp.fe_sample_rate); st != TRANSCRIBE_OK) return st; @@ -351,11 +314,9 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, if (auto st = read_required_f32_kv(gguf, "stt.frontend.f_min", kFamilyTag, hp.fe_f_min); st != TRANSCRIBE_OK) return st; if (auto st = read_required_f32_kv(gguf, "stt.frontend.f_max", kFamilyTag, hp.fe_f_max); st != TRANSCRIBE_OK) return st; - // Prompt MLP. Optional — multilingual variants only - // (nemotron-3.5-asr-streaming-0.6b today). Presence of - // stt.parakeet.prompt.num_prompts (a uint32 KV) is the gate; when - // absent every prompt_* field is left zero/empty and the encoder - // skips the prompt path entirely. + // Prompt MLP. Optional (multilingual variants only). Presence of + // stt.parakeet.prompt.num_prompts is the gate; absent leaves every + // prompt_* field zero/empty and the encoder skips the prompt path. if (gguf_find_key(gguf, "stt.parakeet.prompt.num_prompts") >= 0) { if (auto st = read_required_u32_kv(gguf, "stt.parakeet.prompt.num_prompts", kFamilyTag, hp.prompt_num_prompts); @@ -419,10 +380,7 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, hp.prompt_auto_id = -1; } - // Activation allow-list. The C++ prompt MLP today - // implements ReLU only (matching NeMo's - // EncDecRNNTBPEModelWithPrompt). Other activations are a - // future-variant concern. + // The C++ prompt MLP implements ReLU only. if (hp.prompt_activation != "relu") { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: unsupported prompt activation \"%s\" " @@ -455,9 +413,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } - // Cross-field invariants. These have to hold or the model is - // unbuildable; we catch them here rather than letting them - // surface as confusing shape mismatches downstream. + // Cross-field invariants — caught here rather than surfacing as + // confusing shape mismatches downstream. if (hp.enc_n_layers <= 0 || hp.enc_d_model <= 0 || hp.enc_n_heads <= 0 || hp.enc_d_ff <= 0 || hp.enc_conv_kernel <= 0 || hp.enc_subsampling_factor <= 0 || hp.enc_subsampling_channels <= 0) @@ -480,11 +437,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: joint hparams invalid"); return TRANSCRIBE_ERR_GGUF; } - // Joint activation allow-list. The C++ joint forward implements - // the same three the reference JointNetwork accepts; anything - // else is a converter mistake or a future model we haven't - // ported yet, and producing wrong logits silently is worse than - // failing loudly here. + // Joint activation allow-list (the C++ joint forward implements + // these three). if (hp.joint_activation != "relu" && hp.joint_activation != "sigmoid" && hp.joint_activation != "tanh") @@ -497,12 +451,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, } } if (hp.head_kind == HeadKind::TDT) { - // TDT durations: must be non-empty, must match num_extra_outputs - // (each duration logit corresponds to one durations[i] choice), - // and every entry must be non-negative (negative jumps would - // walk backwards in time and break the decode loop's invariant). - // Zero is allowed and is the standard "stay on this frame, just - // emit a token without advancing" duration. + // TDT durations: non-empty, length == num_extra_outputs, all + // non-negative (negative jumps would walk backwards in time). if (hp.tdt_durations.empty()) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: stt.parakeet.tdt.durations must be non-empty (head_kind=tdt)"); @@ -563,29 +513,18 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, return TRANSCRIBE_ERR_GGUF; } if (hp.fe_type != "mel") { - // Recognized-but-unsupported frontend type. Surfaces as - // ERR_GGUF for now; if we add a non-mel STT family later, - // this turns into a NOT_IMPLEMENTED branch. log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: unsupported frontend type \"%s\" (only \"mel\")", hp.fe_type.c_str()); return TRANSCRIBE_ERR_GGUF; } - // Frontend value-domain checks. The loader reads several string - // and integer fields that the C++ MelFrontend implementation - // (src/transcribe-mel.cpp) does not actually parameterize on. If - // a GGUF carries a value the implementation cannot honor, fail - // here loud and early — letting it through would silently - // substitute the hard-coded behavior at runtime and produce - // wrong features without any error. - // - // When the C++ frontend grows support for more variants (e.g. - // hamming/blackman windows or all_features normalize), the - // corresponding check below should be loosened. - - // Window: only symmetric Hann is implemented - // (build_hann_window_symmetric_padded in transcribe-mel.cpp). + // Frontend value-domain checks: the loader reads fields the C++ + // MelFrontend doesn't parameterize on, so fail loud if a GGUF asks + // for a value the implementation can't honor rather than silently + // substituting the hard-coded behavior. + + // Window: only symmetric Hann is implemented. if (hp.fe_window != "hann") { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: unsupported frontend window \"%s\" " @@ -594,14 +533,9 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, return TRANSCRIBE_ERR_GGUF; } - // Normalize: NeMo's "per_feature" (unbiased per-mel-bin mean/std) - // for offline Parakeet variants, and "none" for streaming variants - // (nemotron-speech-streaming-en-0.6b) whose feature normalisation - // is baked into training rather than applied at inference. NeMo - // also supports "all_features" and CMVN ("fixed_mean"/"fixed_std") - // but no published Parakeet variant uses them, and silently - // substituting per-feature on a GGUF that asked for something else - // would be a confusing bug. + // Normalize: "per_feature" (offline variants) or "none" (streaming, + // normalisation baked into training). NeMo's all_features / CMVN are + // not implemented. if (hp.fe_normalize != "per_feature" && hp.fe_normalize != "none") { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: unsupported frontend normalize \"%s\" " @@ -610,12 +544,8 @@ transcribe_status read_parakeet_hparams(const gguf_context * gguf, return TRANSCRIBE_ERR_GGUF; } - // n_fft must be a power of two: the FFT in transcribe-mel.cpp is - // a radix-2 Cooley-Tukey, and its bit-reversal loop produces - // garbage for non-power-of-two sizes rather than failing - // cleanly. Catch it here so a future converter that emits e.g. - // n_fft=400 fails at load instead of producing nonsense mels. - // (fe_n_fft > 0 was already enforced above.) + // n_fft must be a power of two: the FFT is a radix-2 Cooley-Tukey + // whose bit-reversal produces garbage for non-power-of-two sizes. if ((hp.fe_n_fft & (hp.fe_n_fft - 1)) != 0) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: frontend n_fft (%d) must be a power of 2 " @@ -635,38 +565,20 @@ namespace { using transcribe::weights::lname; -// The canonical find_tensor() + lname() helpers live in -// src/transcribe-weights-util.{h,cpp}; see that header for rationale. -// They are shared between every per-family weights.cpp. The GET_* -// macros below still live here because their type allowlists encode -// a per-family quantization policy, not a shared convention. +// find_tensor() + lname() live in src/transcribe-weights-util.{h,cpp}; +// the GET_* macros below live here because their type allowlists encode +// a per-family quantization policy. Three buckets, each with a fixed +// type allowlist: +// GET_F32 - norms, biases, BN running stats, attn pos_bias (tiny + +// precision-sensitive; fp32 across every preset). +// GET_CONV - conv kernels (TRANSCRIBE_QUANT_CONV_TYPES = F32/F16; +// quantized kernels would need a quantized im2col path +// ggml lacks). +// GET_LIN - mul_mat linear weights (TRANSCRIBE_QUANT_LINEAR_TYPES; +// every family accepts the same set so a quantize preset +// never produces a tensor the loader refuses). constexpr const char * kTag = kFamilyTag; -// Sugar for the GET_OR_FAIL pattern. Variadic so the caller can pass -// the expected dims as a comma-separated list rather than wrapping -// them in extra braces (which would otherwise turn into a GNU -// statement expression when piped through a function-like macro). -// -// Three macros, one per tensor bucket. Each bucket has a fixed -// allowlist of acceptable ggml types; the bucket choice at the call -// site is part of the catalog and never changes per-converter: -// -// GET_F32 - norms, biases, BatchNorm running stats, attn pos_bias. -// These are precision-sensitive and tiny, so they stay -// fp32 across every quant preset the converter ships. -// -// GET_CONV - conv kernels (pre_encode + per-block conv module). -// Uses the project-wide TRANSCRIBE_QUANT_CONV_TYPES -// allowlist (F32 / F16). Quantized kernels would require a -// quantized im2col path that ggml does not provide. -// -// GET_LIN - linear weights consumed by ggml_mul_mat (encoder FF + -// attention projections, predictor LSTM gate matrices, -// joint enc/pred/out projections, predictor embed table). -// Uses the project-wide TRANSCRIBE_QUANT_LINEAR_TYPES -// allowlist. Every family must accept the same set so a -// tools/transcribe-quantize preset never produces a -// tensor the loader refuses. #define GET_F32(slot, name, ...) \ do { \ ggml_tensor * _t = transcribe::weights::find_tensor( \ @@ -717,28 +629,13 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, const int64_t joint_h = hp.joint_hidden; // 0 for CTC const int64_t joint_n = hp.joint_n_classes(); // meaningless for CTC; we read CTC head shape instead - // The flattened "freq" axis going into pre_encode.out is - // (subsampling_channels * F') where F' is the freq dim after the - // three stride-2 convs. The pre-encode convs are k=3, s=2 (NeMo's - // dw_striding stack); their padding tracks NeMo's - // `encoder.causal_downsampling` flag, NOT the conformer's - // conv_context_size (which is the kernel-9 conv-module's context - // and lives in hp.enc_conv_context_{left,right}). The two are - // unrelated. - // - // causal_downsampling=false (every offline variant and the - // buffered-streaming parakeet-unified-en-0.6b): symmetric - // (k-1)/2 padding on both sides. 128 → 64 → 32 → 16; 256*16 - // = 4096. - // causal_downsampling=true (cache-aware streaming variants, - // i.e. nemotron-speech-streaming-en-0.6b): CausalConv2D with - // (left=k-1, right=stride-1) on both axes. 128 → 65 → 33 → - // 17; 256*17 = 4352. - // - // We don't currently emit the boolean to GGUF; we infer it from - // the attention style. Only ChunkedLimited (2-tuple, cache-aware) - // implies causal pre-encode. Regular (offline full / local) and - // ChunkedLimitedWithRc (3-tuple buffered) both use non-causal. + // The flattened "freq" axis into pre_encode.out is + // (subsampling_channels * F'), where F' is the freq dim after three + // k=3 s=2 convs. Padding tracks NeMo's causal_downsampling (inferred + // from the attention style — only ChunkedLimited is causal), NOT the + // conformer conv_context_size: + // non-causal: symmetric (k-1)/2 both sides. 128→64→32→16; 256*16=4096. + // causal: (left=k-1, right=stride-1). 128→65→33→17; 256*17=4352. auto pre_encode_F_prime = [&]() { const bool causal = (hp.enc_att_context_style == @@ -785,11 +682,8 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, GET_LIN(b.ff1_lin1_w, lname("enc.blocks.%d.ff1.linear1.weight", i), d_model, d_ff); GET_LIN(b.ff1_lin2_w, lname("enc.blocks.%d.ff1.linear2.weight", i), d_ff, d_model); - // Self-attention with relative position. Linears can carry - // biases when `enc.use_bias=true` (1.1B / rnnt / ctc variants); - // v2/v3 ship without (Q/K/V/out + pos = bias-free). The - // per-head pos_bias_u/v live alongside regardless. NeMo's - // `linear_pos` is bias-free even when use_bias=true. + // Self-attention with relative position. Q/K/V/out biases only + // when use_bias=true; linear_pos is bias-free even then. GET_F32(b.norm_attn_w, lname("enc.blocks.%d.norm_attn.weight", i), d_model); GET_F32(b.norm_attn_b, lname("enc.blocks.%d.norm_attn.bias", i), d_model); GET_LIN(b.attn_q_w, lname("enc.blocks.%d.attn.linear_q.weight", i), d_model, d_model); @@ -828,11 +722,8 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, GET_F32(b.norm_out_w, lname("enc.blocks.%d.norm_out.weight", i), d_model); GET_F32(b.norm_out_b, lname("enc.blocks.%d.norm_out.bias", i), d_model); - // Optional encoder linear/conv biases. Loaded only when - // hp.enc_use_bias=true (1.1B FastConformer-XL, rnnt-*, ctc-*). - // The 11 names mirror the converter's ENCODER_BLOCK_BIAS_TABLE - // exactly. NeMo's `linear_pos` is bias-free even when - // use_bias=true so attn_pos has no bias slot. + // Optional encoder linear/conv biases (use_bias=true only; + // linear_pos has no bias slot). if (hp.enc_use_bias) { GET_F32(b.ff1_lin1_b, lname("enc.blocks.%d.ff1.linear1.bias", i), d_ff); GET_F32(b.ff1_lin2_b, lname("enc.blocks.%d.ff1.linear2.bias", i), d_model); @@ -849,20 +740,11 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, } if (hp.head_kind == HeadKind::CTC) { - // ----- CTC head ----- - // - // Single 1×1 Conv1d projecting d_model -> vocab+1. NeMo's - // ConvASRDecoder stores it as `decoder.decoder_layers.0`; the - // converter flattens to `head.ctc.weight` / - // `head.ctc.bias`. - // - // Tensor shape in PyTorch is [vocab+1, d_model, 1] - // (out_channels, in_channels, kernel). In ggml ne (fast-to-slow) - // that's [1, d_model, vocab+1]. We don't know `vocab+1` from the - // hparams (CTC GGUFs ship neither predictor.vocab nor - // joint_n_classes), so peek at the raw tensor to capture the - // trailing dim, then run the standard find_tensor with a fully - // specified expected shape so the type+rank checks still fire. + // CTC head: single 1×1 Conv1d projecting d_model -> vocab+1. + // PyTorch shape [vocab+1, d_model, 1] → ggml ne [1, d_model, + // vocab+1]. vocab+1 isn't in the hparams (CTC GGUFs ship neither + // predictor.vocab nor joint_n_classes), so peek at the raw tensor + // for the trailing dim, then run find_tensor with the full shape. ggml_tensor * ctc_peek = ggml_get_tensor(ctx_meta, "head.ctc.weight"); if (ctc_peek == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet: missing tensor \"head.ctc.weight\""); @@ -877,28 +759,19 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, } GET_LIN(weights.ctc_head.weight, "head.ctc.weight", 1, d_model, n_classes); GET_F32(weights.ctc_head.bias, "head.ctc.bias", n_classes); - // Stash the resolved vocab+1 into hp for downstream consumers. - // We're given a const &; we relax that on the way out via a - // const_cast — safe because read_parakeet_hparams / - // build_parakeet_weights are the single producer of hp and its - // only contract is "writable while loading". + // Stash the resolved vocab+1 into hp. const_cast is safe: the + // loader is the single producer of hp, writable while loading. const_cast(hp).head_ctc_n_classes = static_cast(n_classes); } else { // ----- predictor (TDT, RNNT) ----- - // The predictor + joint run on host CPU via decoder.cpp's - // hand-rolled fp32 linear(); the host weight mirror dequantizes - // through ggml_get_type_traits()->to_float on load, so any of the - // GET_LIN-allowed types are safe here. GET_LIN(weights.predictor.embed_w, "pred.embed.weight", pred_h, pred_v); weights.predictor.lstm.assign(hp.pred_n_layers, ParakeetPredictor::LstmLayer{}); for (int i = 0; i < hp.pred_n_layers; ++i) { auto & l = weights.predictor.lstm[i]; - // Both Wx and Wh project from pred_hidden (NeMo's prednet uses - // a learned embedding of size pred_hidden; the embed table - // shape is [pred_vocab, pred_hidden]). The 4*hidden output - // dim is the concatenated (i, f, g, o) gates. + // Wx/Wh both project from pred_hidden; 4*hidden output = + // concatenated (i, f, g, o) gates. const int64_t gates = 4 * pred_h; GET_LIN(l.Wx, lname("pred.lstm.%d.Wx", i), pred_h, gates); GET_LIN(l.Wh, lname("pred.lstm.%d.Wh", i), pred_h, gates); @@ -914,11 +787,8 @@ transcribe_status build_parakeet_weights(ggml_context * ctx_meta, GET_F32(weights.joint.out_b, "joint.out.bias", joint_n); } - // ----- prompt MLP (multilingual variants) ----- - // - // Loaded only when hp.has_prompt is true. The PyTorch nn.Sequential - // indexing is preserved as the canonical name (`.0` input linear, - // `.2` output linear; `.1` is the parameter-free activation). + // ----- prompt MLP (multilingual variants, has_prompt only) ----- + // PyTorch nn.Sequential indexing preserved (.0 input, .2 output). if (hp.has_prompt) { const int64_t prompt_h = hp.prompt_hidden; const int64_t in_dim = static_cast(d_model) + diff --git a/src/arch/parakeet/weights.h b/src/arch/parakeet/weights.h index edb82d36..0b4511a4 100644 --- a/src/arch/parakeet/weights.h +++ b/src/arch/parakeet/weights.h @@ -1,51 +1,27 @@ // arch/parakeet/weights.h - canonical Parakeet tensor catalog and -// per-instance weight slots. +// per-instance weight slots. Internal to src/arch/parakeet/. // -// This header is INTERNAL to src/arch/parakeet/. It defines: +// - ParakeetHParams: architecture KV read from stt.parakeet.* / +// stt.frontend.* before allocating tensors. Every shape-driving dim +// lives here. +// - ParakeetWeights: named borrowed ggml_tensor* slots, one per +// logical weight. Storage is owned by the model's ggml_context; the +// slots must not outlive it. // -// - ParakeetHParams: the architecture KV the loader reads from -// stt.parakeet.* / stt.frontend.* before allocating any tensors. -// Every dim that drives a tensor shape lives here. +// The source file walks the canonical name list and validates each +// tensor (presence + shape) against the hparams. // -// - ParakeetWeights: a struct of named borrowed ggml_tensor* slots, -// one per logical weight in a Parakeet model. The actual storage -// is owned by the model's ggml_context (allocated by -// gguf_init_from_file with no_alloc=false). The slots are -// borrowed pointers and must not outlive the ggml_context. -// -// The corresponding source file walks the canonical tensor name list -// and validates each tensor (presence + shape) against the hparams, -// following llama.cpp's `create_tensor`-with-explicit-shape pattern. -// Required by default; the few honestly-optional tensors are wrapped -// in a NOT_REQUIRED flag. -// -// Naming conventions baked in here (and matched by the converter): -// -// - Linear weight tensors are stored in PyTorch order [out, in]. -// ggml_mul_mat(W, x) takes the activation row vector x and reads -// ne0 of W as the input dim — so PyTorch [out, in] is what we -// want. -// - Conv2d weight tensors are stored in PyTorch OIHW -// [out_channels, in_channels, kH, kW], matching NeMo's native -// layout. Phase 4 figures out whether to feed this directly to -// ggml_conv_2d or transpose again at runtime — that's not 2C's -// problem. -// - Conv1d weight tensors (the depthwise/pointwise convs inside -// the Conformer block) are stored as PyTorch -// [out_channels, in_channels, kernel] for the pointwise convs and -// [out_channels, kernel, 1] for the depthwise (kernel * groups, -// where groups=out_channels=in_channels for depthwise). -// - LSTM gate matrices (Wx, Wh) are stored as a single concatenated -// [4*hidden, input] tensor in PyTorch order (i, f, g, o gates -// stacked along the row dimension). Bias is concatenated [4*hidden]. -// This matches NeMo's safetensors layout exactly so the converter -// does not transpose. -// - LayerNorm and RMSNorm have separate weight + bias tensors of -// shape [d_model]. -// - BatchNorm carries running_mean and running_var alongside -// weight + bias. We keep all four; folding into the surrounding -// conv happens at compute time in phase 4 (encoder is in eval -// mode at inference and BN reduces to an affine transform). +// Tensor layout conventions (matched by the converter): +// - Linear weights: PyTorch [out, in] (ggml_mul_mat reads ne0 = in). +// - Conv2d weights: PyTorch OIHW [out_channels, in_channels, kH, kW]. +// - Conv1d weights: PyTorch [out_channels, in_channels, kernel] for +// pointwise; [out_channels, kernel, 1] for depthwise (groups = +// out_channels = in_channels). +// - LSTM gates (Wx, Wh): single concatenated [4*hidden, input] in +// PyTorch (i, f, g, o) order; bias concatenated [4*hidden]. +// - LayerNorm/RMSNorm: separate weight + bias, shape [d_model]. +// - BatchNorm: weight + bias + running_mean + running_var; all four +// kept (folded into the conv at compute time, eval mode → affine). #pragma once @@ -62,20 +38,13 @@ struct ggml_tensor; namespace transcribe::parakeet { -// --------------------------------------------------------------------------- -// Hyperparameters -// --------------------------------------------------------------------------- -// -// Read from stt.parakeet.* / stt.frontend.* via read_parakeet_hparams(). -// Every field that controls a tensor shape MUST live here so the -// loader can validate weight shapes deterministically against KV -// (rather than against tensor sizes the converter happened to write). - -// Decoder head dispatch. Resolved from stt.parakeet.head_kind (string KV). -// Legacy v2/v3 GGUFs predate the KV; absent KV defaults to TDT, which is -// what those GGUFs are. Drives conditional KV reads in -// read_parakeet_hparams (predictor/joint/tdt only when applicable) and -// the per-head decoder dispatch in model.cpp run(). +// Hyperparameters, read from stt.parakeet.* / stt.frontend.* via +// read_parakeet_hparams(). Every shape-driving field lives here so the +// loader validates weight shapes against KV, not against tensor sizes. + +// Decoder head dispatch (stt.parakeet.head_kind string KV; absent +// defaults to TDT for legacy v2/v3). Drives conditional KV reads and the +// per-head decoder dispatch in run(). enum class HeadKind { TDT, RNNT, CTC }; struct ParakeetHParams { @@ -92,129 +61,84 @@ struct ParakeetHParams { int32_t enc_pos_emb_max_len = 0; bool enc_use_bias = false; // NeMo's RelPositionalEncoding multiplies x by sqrt(d_model) when - // `xscaling=True` is set in the encoder cfg. This pre-block scaling - // is structurally different from how v2/v3/tdt-* (xscaling=False) - // feed the pre_encode output into block 0. Default to false so - // legacy GGUFs without this KV (v2/v3 etc.) keep working unchanged. + // xscaling=True. Default false so legacy GGUFs (v2/v3, xscaling=False) + // keep working unchanged. bool enc_xscaling = false; - // Local attention window. -1 / -1 = full attention (the default for - // most variants). When non-negative, each query attends only to keys - // within an [left, right]-bounded set whose exact semantics depend - // on enc_att_context_style below. - // - // The two styles in use today: - // - // AttContextStyle::Regular — every other variant. When - // both left/right are >= 0 (parakeet-tdt_ctc-1.1b sets - // [128, 128]), pos_emb is shortened to (left+right+1) and the - // band mask sits in matrix_bd via the existing pad/slice path, - // matching NeMo's LocalAttRelPositionalEncoding. - // - // AttContextStyle::ChunkedLimited — nemotron-speech-streaming-en-0.6b - // ([70, 13]). Full RelPositionalEncoding (pos_emb stays at - // 2T-1); a separate chunked mask is built host-side and added - // to matrix_bd. chunk_size = right+1, left_chunks = left/chunk_size. + // Local attention window. -1/-1 = full attention (most variants). + // When non-negative, each query attends to keys within [left, right] + // per enc_att_context_style below: + // Regular — pos_emb shortened to (left+right+1) and the band + // mask in matrix_bd (NeMo LocalAttRelPositionalEncoding); + // parakeet-tdt_ctc-1.1b sets [128, 128]. + // ChunkedLimited — full RelPositionalEncoding (pos_emb 2T-1) plus a + // host-side chunked mask; chunk_size = right+1, left_chunks = + // left/chunk_size. nemotron-speech-streaming-en-0.6b ([70, 13]). int32_t enc_att_context_left = -1; int32_t enc_att_context_right = -1; - // Multi-lookahead training menu for cache-aware streaming models. - // nemotron-speech-streaming-en-0.6b is trained on four (L, R) pairs - // simultaneously (att_context_probs = [0.25, 0.25, 0.25, 0.25] over - // [[70,13],[70,6],[70,1],[70,0]]) and the caller picks at inference - // time via transcribe_parakeet_stream_ext::att_context_right. - // - // Encoded as a flat GGUF int32 array [L0,R0,L1,R1,...] under the KV - // stt.parakeet.encoder.att_context_size_choices. Index 0 is the - // default (max-context / max-accuracy) setting, matching NeMo's - // att_context_size[0] convention; it always equals - // (enc_att_context_left, enc_att_context_right). Empty for offline - // variants — the loader synthesizes a single-element list from the - // scalar fields above when no choices KV is present, so call sites - // can read uniformly. + // Multi-lookahead training menu for cache-aware streaming models. A + // flat GGUF int32 array [L0,R0,L1,R1,...]; index 0 is the default + // (max-context) setting and always equals (enc_att_context_left, + // enc_att_context_right). The caller picks via + // transcribe_parakeet_stream_ext::att_context_right. Empty for offline + // variants; the loader synthesizes a one-element list from the scalar + // fields so call sites read uniformly. nemotron-speech-streaming-en + // is trained on [[70,13],[70,6],[70,1],[70,0]]. std::vector> enc_att_context_size_choices; - // Chunked-limited-with-rc training menu (buffered streaming variants, - // i.e. parakeet-unified-en-0.6b). NeMo stores the 3-tuple [L, C, R] - // training menu as three independent lists in - // `att_chunk_context_size` — the cartesian product of the three is - // the full set of (L, C, R) combinations the model was trained - // against. At inference the runtime picks one entry from each list - // to form the active (L, C, R) tuple driving the chunked attention - // mask. Empty for variants that do not declare - // enc_att_context_style == ChunkedLimitedWithRc. - // - // For parakeet-unified-en-0.6b: L ∈ {70}, C ∈ {1, 2, 7, 13}, - // R ∈ {0, 1, 2, 3, 4, 7, 13} (encoder frames at the 80ms rate). - // The "best accuracy" default is (70, 13, 13) = 5.6s / 1.04s / 1.04s. + // Chunked-limited-with-rc training menu (buffered streaming, i.e. + // parakeet-unified-en-0.6b): three independent L / C / R lists whose + // cartesian product is the set of (L, C, R) tuples the model trained + // against; the runtime picks one entry from each at inference. Empty + // unless enc_att_context_style == ChunkedLimitedWithRc. For + // parakeet-unified-en-0.6b: L ∈ {70}, C ∈ {1,2,7,13}, + // R ∈ {0,1,2,3,4,7,13}; default (70, 13, 13). std::vector enc_att_chunk_left_choices; std::vector enc_att_chunk_chunk_choices; std::vector enc_att_chunk_right_choices; - // Self-attention context style (introduced by streaming variants). - // Resolved from stt.parakeet.encoder.att_context_style (string KV); - // optional, defaults to "regular" so every legacy GGUF stays on the - // existing path. - // - // Regular — full attention (when L=R=-1) or local - // sliding window. Every offline parakeet variant ships this. - // - // ChunkedLimited — cache-aware streaming with a 2-tuple - // (L, R) chunk mask. nemotron-speech-streaming-en-0.6b. - // - // ChunkedLimitedWithRc — buffered streaming with a 3-tuple - // (L, C, R) chunk mask. parakeet-unified-en-0.6b. The training - // menu lives in enc_att_chunk_*_choices above; the active - // tuple is picked at stream_begin time. Note the offline path - // stays on full attention even when the style is - // ChunkedLimitedWithRc — the cfg's `att_context_size` is - // [-1, -1] for this variant; the streaming mask is engaged + // Self-attention context style (stt.parakeet.encoder.att_context_style + // string KV; optional, default "regular"). + // Regular — full attention (L=R=-1) or local window; + // every offline variant. + // ChunkedLimited — cache-aware streaming, 2-tuple (L, R) mask; + // nemotron-speech-streaming-en-0.6b. + // ChunkedLimitedWithRc — buffered streaming, 3-tuple (L, C, R) mask; + // parakeet-unified-en-0.6b. The offline path stays on full + // attention (cfg att_context_size [-1, -1]); the mask is engaged // only by the buffered driver. enum class AttContextStyle { Regular, ChunkedLimited, ChunkedLimitedWithRc }; AttContextStyle enc_att_context_style = AttContextStyle::Regular; - // Depthwise convolution context window inside the Conformer block. - // -1 / -1 = symmetric (kernel-1)/2 on both sides (every offline - // variant). NeMo's `conv_context_size = "causal"` emits - // [kernel-1, 0] (nemotron-speech-streaming-en-0.6b uses kernel=9, - // so [8, 0]) so the depthwise conv only consumes left-context. + // Depthwise conv context window inside the Conformer block. -1/-1 = + // symmetric (kernel-1)/2 (every offline variant). NeMo's causal + // conv_context_size emits [kernel-1, 0] (left-context only). int32_t enc_conv_context_left = -1; int32_t enc_conv_context_right = -1; - // Conv-module normalisation choice. NeMo's offline default is - // BatchNorm; streaming variants use LayerNorm to avoid stat - // brittleness across chunks. With LayerNorm the GGUF carries the - // same `conv.bn.weight` / `conv.bn.bias` tensor names but omits - // running_mean / running_var, and the conformer applies an - // unfused LayerNorm (compute per-channel mean/std at inference, - // then affine-transform with bn_w / bn_b). + // Conv-module normalisation. Offline default BatchNorm; streaming + // variants use LayerNorm. With LayerNorm the GGUF carries the same + // conv.bn.weight / conv.bn.bias names but omits running_mean/var, and + // the conformer applies an unfused LayerNorm (per-channel mean/std + // then affine with bn_w / bn_b). enum class ConvNormType { BatchNorm, LayerNorm }; ConvNormType enc_conv_norm_type = ConvNormType::BatchNorm; - // Cache-aware streaming pre-encode constants. Populated from the - // optional KVs: - // stt.parakeet.encoder.streaming.pre_encode_cache_size (int32) - // stt.parakeet.encoder.streaming.drop_extra_pre_encoded (int32) - // Both are properties of the ConvSubsampling stack — independent of - // the chosen att_context_size — so they're read once and reused - // across all latency settings. - // - // enc_stream_pre_encode_cache_size: number of mel-history frames the - // streaming encoder must prepend to each non-first chunk to fill - // the conv-subsample receptive field (9 on nemotron). - // enc_stream_drop_extra_pre_encoded: number of encoder frames - // discarded after the subsample stack on each non-first chunk to - // align with the cache boundary (2 on nemotron). - // - // Zero when absent from the GGUF (offline variants); the streaming - // code paths gate on both > 0. + // Cache-aware streaming pre-encode constants (optional KVs under + // stt.parakeet.encoder.streaming.*; both 0 for offline variants, on + // which the streaming paths gate). + // pre_encode_cache_size: mel-history frames prepended to each + // non-first chunk to fill the conv-subsample receptive field + // (9 on nemotron). + // drop_extra_pre_encoded: encoder frames discarded after the + // subsample stack per non-first chunk to align with the cache + // boundary (2 on nemotron). int32_t enc_stream_pre_encode_cache_size = 0; int32_t enc_stream_drop_extra_pre_encoded = 0; - // ConvSubsampling's "first-chunk output frames" — used to compute + // ConvSubsampling's first-chunk output frames; used for // chunk_size_first = sampling_frames_first + subsampling_factor * R. - // FastConformer's 2-stride-2 stack ships sampling_frames=[1, 8]; - // we capture the first entry (1). Defaults to subsampling_factor - // when absent (matches the steady-state value, i.e. legacy GGUFs - // collapse to "no first-chunk special case"). + // FastConformer ships sampling_frames=[1, 8] (first entry 1). Defaults + // to subsampling_factor when absent (no first-chunk special case). int32_t enc_stream_sampling_frames_first = 0; // Predictor (RNN-T prediction network). TDT and RNNT only; zero for CTC. @@ -226,44 +150,28 @@ struct ParakeetHParams { // Joint network. TDT and RNNT only; zero for CTC. int32_t joint_hidden = 0; int32_t joint_num_extra_outputs = 0; // TDT durations count; 0 for RNNT - // Joint output activation. NeMo Parakeet 0.6B v2/v3 ship "relu" - // (verified against config.json on both variants); the rnnt.py - // reference defaults to "tanh" but no published Parakeet variant - // we know about uses tanh. Read from stt.parakeet.joint.activation - // and validated against the small allow-list the C++ joint forward - // implements. Without this read at load time the C++ side would - // hard-code an activation choice and silently produce wrong logits - // on a model that asked for a different one. + // Joint output activation (stt.parakeet.joint.activation), validated + // against the C++ joint forward's allow-list at load. v2/v3 ship + // "relu". Read so the C++ side doesn't hard-code an activation and + // silently produce wrong logits on a model that asked for another. std::string joint_activation; - // TDT decoding parameters. - // - // tdt_durations: the duration values the joint network's "extra - // outputs" axis predicts (e.g. [0,1,2,3,4] for the published 0.6B - // variants). The argmax over duration logits selects an index into - // this list; the chosen integer is how many encoder frames the TDT - // decode driver advances after emitting a token. PLAN.md fixes the - // KV name as `stt.parakeet.tdt.durations` (i32 array). Length must - // equal joint_num_extra_outputs (we cross-validate at load time). + // TDT durations: the values the joint's "extra outputs" axis predicts + // (e.g. [0,1,2,3,4]); argmax over duration logits indexes this list, + // and the chosen integer is how many encoder frames the decoder + // advances after a token. KV stt.parakeet.tdt.durations; length must + // equal joint_num_extra_outputs (cross-validated at load). std::vector tdt_durations; - // tdt_max_symbols: the per-encoder-frame "stuck" cap from NeMo's - // greedy_batch decoder. Optional KV with a documented default of 10 - // (matches both published v2 and v3 configs). When the decoder has - // emitted N consecutive zero-duration tokens for the same frame - // without advancing time, it forces a +1 advance to break the - // loop. Set to 0 to disable. + // tdt_max_symbols: per-frame "stuck" cap (NeMo greedy_batch). Optional + // KV, default 10. After N consecutive zero-duration tokens on the same + // frame, forces a +1 advance. 0 disables. int32_t tdt_max_symbols = 10; - // Frontend (mel feature extractor). The full set of stt.frontend.* - // KV the converter emits and the phase 3 C++ frontend reads. PLAN.md - // declares this list as the *complete* set the converter and loader - // must agree on, so the loader is the gate that catches any future - // converter that drops a field. Fields irrelevant to Parakeet - // (LFR window/shift, CMVN mean/inv_stddev) are not in this struct - // because no published Parakeet variant uses them — when SenseVoice - // or another LFR/CMVN family lands, those fields move to a - // family-agnostic FrontendParams struct. + // Frontend (mel feature extractor). The complete stt.frontend.* set + // the converter emits and the C++ frontend reads; the loader gates on + // it. CMVN/LFR fields are omitted (no published Parakeet variant uses + // them). std::string fe_type; // "mel" for parakeet int32_t fe_num_mels = 0; int32_t fe_sample_rate = 0; @@ -283,20 +191,15 @@ struct ParakeetHParams { // Zero for TDT/RNNT. int32_t head_ctc_n_classes = 0; - // Language-conditioning prompt MLP (NeMo's - // EncDecRNNTBPEModelWithPrompt). Multilingual variants prepend a - // 2-layer MLP that mixes a one-hot prompt vector into the encoder - // output before the RNN-T joint sees it: - // - // x = concat(enc[d_model], one_hot(prompt_id)[num_prompts]) - // h = relu(W0 @ x + b0) // W0: [prompt_hidden, d_model+P] - // y = W2 @ h + b2 // W2: [d_model, prompt_hidden] - // enc_out := y - // - // Gated on `has_prompt` (true iff stt.parakeet.prompt.num_prompts - // is present in the GGUF). Today: nemotron-3.5-asr-streaming-0.6b. - // Loader reads stt.parakeet.prompt.{num_prompts, hidden, field, - // activation, dictionary.locales, dictionary.indices, auto_id}. + // Language-conditioning prompt MLP (NeMo + // EncDecRNNTBPEModelWithPrompt). Multilingual variants mix a one-hot + // prompt vector into the encoder output before the RNN-T joint: + // x = concat(enc[d_model], one_hot(prompt_id)[num_prompts]) + // h = relu(W0 @ x + b0) // W0: [prompt_hidden, d_model+P] + // y = W2 @ h + b2 // W2: [d_model, prompt_hidden] + // enc_out := y + // Gated on has_prompt (stt.parakeet.prompt.num_prompts present). + // Today: nemotron-3.5-asr-streaming-0.6b. bool has_prompt = false; int32_t prompt_num_prompts = 0; int32_t prompt_hidden = 0; @@ -312,29 +215,18 @@ struct ParakeetHParams { }; // Read every required stt.parakeet.* / stt.frontend.* KV into hp. -// Returns: -// TRANSCRIBE_OK on success. -// TRANSCRIBE_ERR_INVALID_ARG if gguf is null. -// TRANSCRIBE_ERR_GGUF if a required key is missing or has the -// wrong type, or if cross-field invariants -// don't hold (e.g. d_model not divisible -// by n_heads). +// Returns INVALID_ARG if gguf is null, ERR_GGUF on a missing/wrong-type +// key or a cross-field invariant failure (e.g. d_model % n_heads != 0). transcribe_status read_parakeet_hparams(const gguf_context * gguf, ParakeetHParams & hp); -// --------------------------------------------------------------------------- -// Weight slots -// --------------------------------------------------------------------------- -// -// Every ggml_tensor pointer below is a borrowed pointer into the -// model's ggml_context (the one that owns the data buffer). They are -// invalidated when the model is freed. +// Weight slots. Every ggml_tensor* below is borrowed into the model's +// ggml_context and is invalidated when the model is freed. struct ParakeetPreEncode { // dw_striding subsampling: 1 standard conv + 2 depthwise-separable - // conv pairs = 8x downsampling on the time axis. The numeric - // indices in the names below come from NeMo's torch.nn.Sequential - // layout (1 and 4 are activation modules, hence skipped). + // conv pairs = 8x time downsampling. Indices follow NeMo's + // torch.nn.Sequential (1 and 4 are activation modules, skipped). ggml_tensor * conv0_w = nullptr; // [out=channels, in=1, kh=3, kw=3] ggml_tensor * conv0_b = nullptr; // [channels] ggml_tensor * conv2_w = nullptr; // [out=channels, in=1 (dw), kh=3, kw=3] @@ -350,12 +242,9 @@ struct ParakeetPreEncode { ggml_tensor * out_b = nullptr; }; -// One Conformer block. Same shape contract for every block; the loader -// builds n_layers of these. The bias slots (`*_b` on the linear/conv -// layers, NOT on the layer norms / BN / pos_bias) are populated only -// when `enc_use_bias=true` — left null for the v2/v3 baseline. The -// shared `transcribe::conformer::BlockView` already accepts nullable -// biases on every linear/conv slot. +// One Conformer block (same shape contract for every block). The bias +// slots (*_b on linear/conv layers, NOT on norms / BN / pos_bias) are +// populated only when enc_use_bias=true; null for the v2/v3 baseline. struct ParakeetBlock { // Macaron feed-forward 1. ggml_tensor * norm_ff1_w = nullptr; // [d_model] @@ -415,12 +304,11 @@ struct ParakeetBlock { }; struct ParakeetPredictor { - // Token embedding. The +1 row is the "start of sequence / no - // previous token" embedding NeMo prepends to vocab_size. + // Token embedding; the +1 row is NeMo's "start of sequence" embedding. ggml_tensor * embed_w = nullptr; // [pred_vocab, pred_hidden] - // 2-layer LSTM by default for v2/v3 0.6b. Concatenated gate - // weights in PyTorch (i, f, g, o) order, 4*pred_hidden rows. + // 2-layer LSTM (v2/v3 0.6b). Concatenated gate weights in PyTorch + // (i, f, g, o) order, 4*pred_hidden rows. struct LstmLayer { ggml_tensor * Wx = nullptr; // [4*pred_hidden, pred_hidden] (input-to-hidden, layer 0 input dim is pred_hidden because of the embed) ggml_tensor * Wh = nullptr; // [4*pred_hidden, pred_hidden] (hidden-to-hidden) @@ -466,21 +354,12 @@ struct ParakeetWeights { ParakeetPromptMlp prompt; // populated when hp.has_prompt }; -// Walk the canonical tensor list, look up each tensor by name in the -// gguf, validate its shape against hp, and store the borrowed pointer -// in the matching slot of weights. The ggml_tensor objects must -// already exist in ctx_meta — typically because the caller invoked -// gguf_init_from_file with no_alloc=false and ctx=&ctx_meta, which -// both creates ggml_tensors AND reads the data section. -// -// Follows llama.cpp's load_tensors pattern: every required tensor -// must be present with the exact expected shape; missing or -// shape-mismatched tensors return TRANSCRIBE_ERR_GGUF with a log -// message naming the offending tensor. -// -// On failure the partially-built `weights` is left in an -// indeterminate state — the caller is expected to throw the whole -// model away. +// Walk the canonical tensor list, look up each tensor by name, validate +// its shape against hp, and store the borrowed pointer in the matching +// slot. The ggml_tensors must already exist in ctx_meta. Every required +// tensor must be present with the exact shape; missing/mismatched returns +// TRANSCRIBE_ERR_GGUF naming the tensor. On failure `weights` is left +// indeterminate (the caller throws the model away). transcribe_status build_parakeet_weights(ggml_context * ctx_meta, const ParakeetHParams & hp, ParakeetWeights & weights); diff --git a/src/arch/qwen3_asr/capabilities.cpp b/src/arch/qwen3_asr/capabilities.cpp index 8f8bcd8d..eb9c9c25 100644 --- a/src/arch/qwen3_asr/capabilities.cpp +++ b/src/arch/qwen3_asr/capabilities.cpp @@ -9,31 +9,16 @@ void apply_family_invariants(transcribe_model & model) { caps.native_sample_rate = 16000; - // Qwen3-ASR emits transcript text inside a Qwen-style chat response. - // It does not surface timestamps of any kind for the transcript-only - // task. The sibling Qwen3-ForcedAligner produces word timestamps, - // but that is a separate checkpoint with a different head contract - // and would register as its own family (qwen3_forced_aligner) when - // ported. + // Transcript-only chat response; no timestamps (the sibling + // Qwen3-ForcedAligner would be its own family). caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; - // The publisher documents language detection (the LM prefixes its - // transcript with "language X") and also accepts a caller-supplied - // language hint — the chat template seeds the assistant turn with - // "language {Name}" and the LM continues with pure - // transcript text, skipping its own detection preamble. Either - // path is supported; caps.languages is the list of codes both - // support, with the same BCP-47 canonicalization used by every - // other family in this library. - // - // Translation is not an advertised capability for Qwen3-ASR; - // callers get transcript-only output in the audio's source - // language. + // Supports language auto-detect and a caller-supplied hint (caps.languages + // lists the BCP-47 codes). Translation is not advertised. caps.supports_translate = false; - // Cancellation is wired at the per-run level. Whisper-specific - // long-form / temperature-fallback / initial-prompt features do - // not apply here. No runtime PNC/ITN toggle. + // Cancellation is wired at the per-run level. No PNC/ITN toggle; the + // Whisper-specific features do not apply here. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); } diff --git a/src/arch/qwen3_asr/decoder.cpp b/src/arch/qwen3_asr/decoder.cpp index cd13ade3..ecc11186 100644 --- a/src/arch/qwen3_asr/decoder.cpp +++ b/src/arch/qwen3_asr/decoder.cpp @@ -2,20 +2,12 @@ // // Reference: Qwen3ASRThinkerTextModel in modeling_qwen3_asr.py. // -// The per-block math (pre-LN RMSNorm, GQA with per-head Q/K-RMSNorm, -// NeoX RoPE @ θ=1e6, KV write/read, SwiGLU on packed gate_up) is -// shared with arch/funasr_nano via `src/causal_lm/`. This file owns: -// - graph allocation -// - prompt + audio injection (3-way concat: prefix | enc_out | suffix) -// - tensor naming and dump-point preservation for validate.py parity -// - final RMSNorm + tied lm_head head -// - block_prefill / block_step driver loops -// -// MRoPE reduction: for text-only ASR every position has a single -// (T, H, W) coordinate so the interleaved MRoPE collapses to NeoX -// RoPE on the temporal axis. The shared block helpers use NEOX RoPE. -// If a future variant breaks this assumption we'll switch to -// ggml_rope_multi. +// The per-block math (pre-LN RMSNorm, GQA with per-head Q/K-RMSNorm, NeoX +// RoPE @ θ=1e6, KV write/read, SwiGLU on packed gate_up) is shared with +// arch/funasr_nano via `src/causal_lm/`. This file owns graph allocation, +// prompt + audio injection (3-way concat: prefix | enc_out | suffix), tensor +// naming + dump points, the final RMSNorm + tied lm_head, and the +// block_prefill / block_step driver loops. MRoPE reduction: see decoder.h. #include "decoder.h" @@ -121,7 +113,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, const auto block_params = to_block_params(hp); const float rms_eps = hp.dec_rms_norm_eps; - // ---------- Graph inputs ---------- + // Graph inputs. pb.input_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T_prompt); named(pb.input_ids_in, "dec.input_ids"); ggml_set_input(pb.input_ids_in); @@ -148,14 +140,14 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, } pb.graph = gf; - // ---------- Token embedding (all positions, for the dump) ---------- + // Token embedding (all positions, for the dump). ggml_tensor * token_emb_all = ggml_get_rows( ctx, weights.dec_embed.token_w, pb.input_ids_in); named(token_emb_all, "dec.token_emb"); pb.dumps.token_emb = token_emb_all; transcribe::debug::mark_tensor_for_dump(token_emb_all); - // ---------- Audio injection via 3-way concat ---------- + // Audio injection via 3-way concat. // Single contiguous audio block at positions [prefix_len, prefix_len+T_enc). const size_t emb_elem = ggml_element_size(token_emb_all); ggml_tensor * x_prefix = nullptr; @@ -186,7 +178,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, pb.dumps.audio_injected = x; transcribe::debug::mark_tensor_for_dump(x); - // ---------- Block stack ---------- + // Block stack. for (int il = 0; il < n_layer; ++il) { causal_lm::BlockOpts opts {}; opts.use_flash = use_flash; @@ -218,7 +210,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, } } - // ---------- Final RMSNorm + tied lm_head ---------- + // Final RMSNorm + tied lm_head. x = ggml_mul(ctx, ggml_rms_norm(ctx, x, rms_eps), weights.dec_final.norm_w); named(x, "dec.out_before_head"); @@ -259,9 +251,7 @@ PrefillBuild build_prefill_graph(ggml_context * ctx, return pb; } -// --------------------------------------------------------------------------- // Step graph (single-token decode) -// --------------------------------------------------------------------------- StepBuild build_step_graph(ggml_context * ctx, const QwenAsrWeights & weights, @@ -295,7 +285,7 @@ StepBuild build_step_graph(ggml_context * ctx, const auto block_params = to_block_params(hp); - // ---------- Graph inputs (persist across steps) ---------- + // Graph inputs (persist across steps). sb.input_id_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_name(sb.input_id_in, "step.input_id"); ggml_set_input(sb.input_id_in); @@ -323,11 +313,11 @@ StepBuild build_step_graph(ggml_context * ctx, } sb.graph = gf; - // ---------- Token embedding ---------- + // Token embedding. ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed.token_w, sb.input_id_in); // [hidden, 1] - // ---------- Block stack ---------- + // Block stack. for (int il = 0; il < n_layer; ++il) { x = causal_lm::block_step( ctx, gf, x, @@ -342,7 +332,7 @@ StepBuild build_step_graph(ggml_context * ctx, use_flash); } - // ---------- Final RMSNorm + tied lm_head ---------- + // Final RMSNorm + tied lm_head. x = ggml_mul(ctx, ggml_rms_norm(ctx, x, rms_eps), weights.dec_final.norm_w); ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.token_w, x); @@ -416,10 +406,7 @@ PrefillBuildBatched build_prefill_graph_batched( ggml_set_name(pb.input_ids_in, "prefill.input_ids"); ggml_set_input(pb.input_ids_in); - // Audio injection is an elementwise blend over the flat token axis (no - // set_rows): a k-quant token_embd get_rows is unsupported on CUDA and runs - // on CPU; a set_rows consuming that CPU tensor straddles the CPU/CUDA split - // and faults. x*keep_mask + audio_dense crosses the split cleanly. + // Audio injection via elementwise blend (no set_rows); see decoder.h. pb.audio_dense_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, static_cast(T_prompt_max) * B); ggml_set_name(pb.audio_dense_in, "prefill.audio_dense"); @@ -460,10 +447,8 @@ PrefillBuildBatched build_prefill_graph_batched( ggml_tensor * ids_flat = ggml_reshape_1d(ctx, pb.input_ids_in, static_cast(T_prompt_max) * B); - // x is 2D [hidden, T_prompt_max*B] (contiguous). Audio injection: blend the - // enc_out embeds in elementwise, x = x*keep_mask + audio_dense, where - // keep_mask broadcasts over hidden (ne[0]). Then reshape to 3D for the - // batched block stack. + // x = x*keep_mask + audio_dense (keep_mask broadcasts over hidden), then + // reshape to 3D for the batched block stack. ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed.token_w, ids_flat); x = ggml_add(ctx, ggml_mul(ctx, x, pb.keep_mask_in), pb.audio_dense_in); x = ggml_reshape_3d(ctx, x, hidden, T_prompt_max, B); @@ -553,7 +538,7 @@ StepBuildBatched build_step_graph_batched( const auto block_params = to_block_params(hp); - // ---------- Graph inputs (persist across steps) ---------- + // Graph inputs (persist across steps). sb.input_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, B); ggml_set_name(sb.input_ids_in, "step.input_ids"); ggml_set_input(sb.input_ids_in); @@ -578,11 +563,11 @@ StepBuildBatched build_step_graph_batched( } sb.graph = gf; - // ---------- Token embedding ---------- + // Token embedding. ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed.token_w, sb.input_ids_in); // [hidden, B] - // ---------- Block stack ---------- + // Block stack. for (int il = 0; il < n_layer; ++il) { x = causal_lm::block_step_batched( ctx, gf, x, @@ -604,7 +589,7 @@ StepBuildBatched build_step_graph_batched( } } - // ---------- Final RMSNorm + tied lm_head ---------- + // Final RMSNorm + tied lm_head. x = ggml_mul(ctx, ggml_rms_norm(ctx, x, rms_eps), weights.dec_final.norm_w); // [hidden, B] ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.token_w, x); diff --git a/src/arch/qwen3_asr/decoder.h b/src/arch/qwen3_asr/decoder.h index 48e8be3d..44800448 100644 --- a/src/arch/qwen3_asr/decoder.h +++ b/src/arch/qwen3_asr/decoder.h @@ -17,9 +17,17 @@ // multi-modal positions we'll switch to ggml_rope_multi. // // Audio injection: input_ids contains `audio_token_id` placeholders at -// the positions the encoder's output frames should land. The prefill -// graph embeds input_ids as usual, then `ggml_set_rows` scatters the -// encoder output over the audio positions. +// the positions the encoder's output frames should land. Two injection +// mechanisms, neither uses ggml_set_rows: +// - single-utterance prefill (build_prefill_graph): a 3-way ggml_concat +// of [prefix_emb | encoder_output | suffix_emb] along the time axis, +// assuming one contiguous audio block at [prefix_len, prefix_len+T_enc). +// - batched prefill (build_prefill_graph_batched): an elementwise blend +// x*keep_mask + audio_dense, where audio_dense holds the encoder embeds +// scattered host-side into their prompt positions (zero elsewhere) and +// keep_mask is 0 at audio positions / 1 elsewhere. Elementwise ops +// cross the CPU/CUDA split (forced by k-quant token_embd get_rows) +// cleanly, unlike a set_rows. #pragma once @@ -48,7 +56,7 @@ struct PrefillBuild { ggml_tensor * input_ids_in = nullptr; // [T_prompt] i32 (for dec.token_emb dump) ggml_tensor * enc_out_in = nullptr; // [enc_output_dim, T_enc] f32 ggml_tensor * positions_in = nullptr; // [T_prompt] i32 for RoPE - ggml_tensor * mask_in = nullptr; // [T_prompt, T_prompt] f32 (causal) + ggml_tensor * mask_in = nullptr; // [T_prompt, T_prompt] f16 (causal) ggml_tensor * out = nullptr; // [vocab_size] — last-position logits DecoderDumps dumps {}; ggml_cgraph * graph = nullptr; @@ -59,26 +67,13 @@ struct PrefillBuild { int suffix_len = 0; // # prompt tokens after the audio block }; -// Build a prefill graph that: -// 1. fetches token embeddings for the full prompt (for dump parity), -// 2. builds the injected-embedding sequence by concatenating -// [prefix_token_embeddings | encoder_output | suffix_token_embeddings], -// 3. runs the 28 Qwen3 blocks, writing K/V into kv_cache at positions -// [0, T_prompt), -// 4. applies the final RMSNorm and tied-head projection, -// 5. outputs logits for the last position only. -// -// First port assumption: the audio block is a single contiguous run at -// positions [prefix_len, prefix_len + T_enc). Multi-audio prompts (not -// used by Qwen3-ASR's chat template) would need repeated get_rows + -// concat segments. -// -// The kv_cache is WRITTEN by the graph. Callers should set -// kv_cache.n = T_prompt and kv_cache.head = T_prompt after compute. -// kv_batch_slot / kv_n_batch: for offline batched decode, write this -// utterance's KV into slab `kv_batch_slot` of a batched cache holding -// `kv_n_batch` slabs (see causal_lm::block_prefill / kv_init_batched). -// Defaults (0, 1) reproduce the single-shot KV layout exactly. +// Build a prefill graph: token-embed the full prompt, concat +// [prefix_emb | encoder_output | suffix_emb], run the Qwen3 blocks (writing +// K/V into kv_cache at [0, T_prompt)), final RMSNorm + tied head, output +// last-position logits. Assumes a single contiguous audio block at +// [prefix_len, prefix_len + T_enc). Callers set kv_cache.n / .head = T_prompt +// after compute. kv_batch_slot / kv_n_batch route this utterance's KV into a +// batched cache slab; defaults (0, 1) reproduce the single-shot layout. PrefillBuild build_prefill_graph(ggml_context * ctx, const QwenAsrWeights & weights, const QwenAsrHParams & hp, @@ -105,18 +100,14 @@ struct StepBuild { int max_n_kv = 0; // static shape sized for whole run }; -// Build a static-shape single-token step graph suitable for reuse -// across every autoregressive step in a run. Topology depends only on -// max_n_kv (the max KV window the LM will see), not on n_past. The -// graph has four input tensors the caller updates per step: +// Build a static-shape single-token step graph reused across every step. +// Topology depends only on max_n_kv (the max KV window), not n_past. Four +// per-step input tensors: // input_id_in [1] — the token to embed // position_in [1] — RoPE position // kv_idx_in [1] — where to write K/V (via ggml_set_rows) // mask_in [max_n_kv, 1] — attention mask; positions > n_past hold -inf -// -// Reusing the same graph every step eliminates the ggml_free/init + -// rebuild + sched_reset + sched_alloc_graph overhead that per-step -// construction incurs (measured ~0.4 ms/step before this change). +// Reuse avoids per-step rebuild + sched_alloc overhead (~0.4 ms/step). StepBuild build_step_graph(ggml_context * ctx, const QwenAsrWeights & weights, const QwenAsrHParams & hp, @@ -129,10 +120,10 @@ StepBuild build_step_graph(ggml_context * ctx, struct PrefillBuildBatched { ggml_tensor * input_ids_in = nullptr; // [T_prompt_max, B] i32 (audio_token placeholders) // Audio injection via elementwise blend (no set_rows): audio_dense holds the - // enc_out embeds scattered (host-side) into their prompt positions, zero - // elsewhere; keep_mask is 0 at audio positions and 1 elsewhere. The block - // input is x*keep_mask + audio_dense. Elementwise ops cross the CPU/CUDA - // split (forced by k-quant token_embd get_rows) cleanly, unlike a set_rows. + // enc_out embeds scattered host-side into their prompt positions (zero + // elsewhere), keep_mask is 0 there and 1 elsewhere, block input is + // x*keep_mask + audio_dense. Elementwise ops cross the CPU/CUDA split + // (forced by k-quant token_embd get_rows) cleanly, unlike a set_rows. ggml_tensor * audio_dense_in = nullptr; // [dec_hidden, T_prompt_max*B] f32 ggml_tensor * keep_mask_in = nullptr; // [1, T_prompt_max*B] f32 (0=audio,1=keep) ggml_tensor * positions_in = nullptr; // [T_prompt_max] i32 (shared 0..T-1) @@ -150,10 +141,9 @@ struct PrefillBuildBatched { // Build a batched prefill graph: B prompts of up to T_prompt_max tokens each, // writing each utterance's K/V into slab b of a batched cache (kv_init_batched). -// Audio is injected by ggml_set_rows scatter of enc_out into the audio -// placeholder positions (per-utterance scatter indices in enc_idx_in). The -// last-real-position logits per utterance are gathered via last_idx_in, so the -// readback is [vocab, B] (not [vocab, T_prompt_max, B]). Requires use_flash. +// Audio is injected by the elementwise blend (audio_dense_in / keep_mask_in). +// Per-utterance last-real-position logits are gathered via last_idx_in, so the +// readback is [vocab, B]. Requires use_flash. PrefillBuildBatched build_prefill_graph_batched( ggml_context * ctx, const QwenAsrWeights & weights, @@ -179,9 +169,8 @@ struct StepBuildBatched { int n_batch = 0; }; -// Build a static-shape batched step graph reused across every step of an -// offline batch of `n_batch` utterances against a batched KV cache -// (causal_lm::kv_init_batched). The four input tensors are updated per step: +// Build a static-shape batched step graph reused across every step for an +// offline batch of `n_batch` utterances. Four per-step input tensors: // input_ids_in [B] — the token to embed for each utterance // position_in [B] — RoPE position per utterance (= its n_past) // kv_idx_in [1, B] — KV write row per utterance (= its n_past) diff --git a/src/arch/qwen3_asr/encoder.cpp b/src/arch/qwen3_asr/encoder.cpp index d3657658..044b175e 100644 --- a/src/arch/qwen3_asr/encoder.cpp +++ b/src/arch/qwen3_asr/encoder.cpp @@ -18,9 +18,7 @@ namespace transcribe::qwen3_asr { -// --------------------------------------------------------------------------- // Timing / cu_seqlens -// --------------------------------------------------------------------------- int32_t aftercnn_len(int32_t mel_len) { // Three rounds of Conv2d(stride=2, pad=1, kernel=3): @@ -91,9 +89,7 @@ std::vector build_cu_seqlens_mask(const EncoderTiming & t, return std::vector(static_cast(T) * T, 0.0f); } -// --------------------------------------------------------------------------- // Graph construction -// --------------------------------------------------------------------------- namespace { diff --git a/src/arch/qwen3_asr/encoder.h b/src/arch/qwen3_asr/encoder.h index 43429fa4..a3d26e15 100644 --- a/src/arch/qwen3_asr/encoder.h +++ b/src/arch/qwen3_asr/encoder.h @@ -5,13 +5,9 @@ // qwen_asr.core.transformers_backend.modeling_qwen3_asr // .Qwen3ASRAudioEncoder. // -// Attention note: the reference's transformers/eager path ignores -// `cu_seqlens` and runs full bidirectional attention over the -// pad-select-trimmed sequence. We match that (see -// build_cu_seqlens_mask). vLLM's flash-attn-2 path honors -// `cu_seqlens` to chunk at `n_window_infer`, but measurements show -// the eager path has better LibriSpeech WER and that's the reference -// we validate against. +// 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. // // Shape conventions (ggml fast-to-slow ne[]): // @@ -71,12 +67,9 @@ EncoderTiming compute_encoder_timing(int32_t n_mel_frames, 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. Named -// `build_cu_seqlens_mask` for history; the reference's `cu_seqlens` -// block-diagonal pattern is only applied by vLLM's flash-attn-2 -// path, and measurements show the eager full-attention baseline has -// better WER on LibriSpeech. +// 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); diff --git a/src/arch/qwen3_asr/model.cpp b/src/arch/qwen3_asr/model.cpp index 15c7d400..cd1e632d 100644 --- a/src/arch/qwen3_asr/model.cpp +++ b/src/arch/qwen3_asr/model.cpp @@ -1,11 +1,4 @@ // arch/qwen3_asr/model.cpp - Qwen3-ASR family handler. -// -// Wiring-only skeleton. `load()` is real — it reads the GGUF KV, -// resolves the tokenizer + frontend, and builds the tensor catalog — -// so that GGUF conversion round-trips can be verified before inference -// code lands. `init_context()` and `run()` currently return -// TRANSCRIBE_ERR_NOT_IMPLEMENTED and will be filled in once the -// encoder and LM graph builders exist. #include "qwen3_asr.h" @@ -86,19 +79,14 @@ namespace { constexpr const char k_default_variant[] = "qwen3-asr"; -// --------------------------------------------------------------------------- // Input-length contract (see docs/input-limits.md). Qwen3-ASR is a // hard-context-cap family: audio tokens + chat prompt + generation share the -// Qwen3 decoder's context window (dec_max_position_embeddings). The KV cache -// grows to fit the prompt, clamped to that ceiling (replacing the earlier -// fixed 2048 wall that ignored the model's real context). Over-length input -// is rejected up front with TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that -// fills the per-run generation budget before end-of-stream is flagged via -// transcribe_was_truncated(). -// --------------------------------------------------------------------------- - -// Per-run generation budget (matches the reference dumper default; the -// decode loop is intentionally left unchanged — see the family doc). +// Qwen3 decoder's context window (dec_max_position_embeddings), clamped to that +// ceiling. Over-length input is rejected with 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 (matches the reference dumper default). constexpr int k_max_new = 256; // Effective decoder context ceiling, in tokens: the model's trained maximum, @@ -165,30 +153,25 @@ transcribe_status load( if (const transcribe_status st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) return st; - // Chat template (optional, but required at run time — we read it - // here so the absence surfaces at load rather than mid-decode). + // Chat template (read here so its absence surfaces at load, not mid-decode). (void)read_optional_string_kv( loader.gguf(), "tokenizer.chat_template", "qwen3_asr", "", m->chat_template); - // Resolve the chat-template special-token ids at load so a vocab - // drift (e.g. a future fine-tune that renames or reorders the - // role tokens) surfaces here instead of silently producing a - // wrong prompt at decode time. + // Resolve chat-template special-token ids at load so a vocab drift surfaces + // here instead of silently producing a wrong prompt at decode time. if (const transcribe_status st = resolve_chat_tokens(m->tok, m->chat_tokens); st != TRANSCRIBE_OK) return st; - // Hparams. if (const transcribe_status st = read_qwen3_asr_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) return st; // Publish the input-length ceiling now that the decoder context window - // and frontend rate are known. See docs/input-limits.md. + // and frontend rate are known. m->caps.max_audio_ms = qwen3_max_audio_ms(m->hparams); - // Basis for the session-level limits query (transcribe_session_get_limits): - // the same constants qwen3_max_audio_ms uses, kept so the effective limit - // can be recomputed at a lowered session n_ctx. + // Basis for transcribe_session_get_limits: the same constants + // qwen3_max_audio_ms uses, so the limit recomputes at a lowered n_ctx. if (m->hparams.dec_max_position_embeddings > 0 && m->hparams.fe_hop_length > 0 && m->hparams.fe_sample_rate > 0) { m->limits.has_context_cap = true; @@ -234,11 +217,8 @@ transcribe_status load( cfg.window_type = m->hparams.fe_window; // "hann_periodic" cfg.normalize = m->hparams.fe_normalize; // "per_utterance" (Whisper) - // Optional filterbank + window buffers baked by the converter - // from librosa / Whisper. If present, MelFrontend uses them - // instead of reconstructing from hparams. Removes filterbank/ - // window as a variable during numerical bring-up. Absent falls - // back to the in-code builders. + // Optional filterbank + window buffers baked by the converter; if + // present MelFrontend uses them instead of reconstructing from hparams. { using R = transcribe::load_common::ReadF32Result; const size_t fb_elems = @@ -264,7 +244,7 @@ transcribe_status load( m->mel.emplace(cfg); } - // Stage 2: reopen with no_alloc to build the tensor catalog. + // Reopen with no_alloc to build the tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -391,15 +371,10 @@ transcribe_status init_context( return TRANSCRIBE_OK; } -// Resolve the narrow set of chat-template piece strings against the -// loaded tokenizer, filling `out`. Hard-fails with TRANSCRIBE_ERR_GGUF -// if any piece is missing — a future checkpoint that reorders the -// vocab will surface here at load time, before any prompt is built. -// -// The newline piece is stored under its GPT-2 byte-level-encoded form -// (\n = 0x0A → Unicode codepoint U+010A "Ċ", UTF-8 \xC4\x8A), same as -// every other non-printable byte in the Qwen3 vocab. The role names -// and the two <|im_*|> special tokens are stored verbatim. +// Resolve the chat-template piece strings against the loaded tokenizer. +// Hard-fails with TRANSCRIBE_ERR_GGUF on a missing piece (a vocab reorder +// surfaces here at load). The newline is stored in its GPT-2 byte-level form +// (\n → U+010A "Ċ", \xC4\x8A); roles and <|im_*|> tokens are verbatim. transcribe_status resolve_chat_tokens(const transcribe::Tokenizer & tok, ChatTokens & out) { @@ -428,21 +403,16 @@ transcribe_status resolve_chat_tokens(const transcribe::Tokenizer & tok, return TRANSCRIBE_OK; } -// Build the prompt token sequence + audio-position list for a single -// utterance. Mirrors the Qwen3-ASR chat template at the token level: +// Build the prompt token sequence + audio-position list, mirroring the +// Qwen3-ASR chat template at the token level: // // <|im_start|>system\n<|im_end|>\n // <|im_start|>user\n<|audio_start|><|audio_pad|>*T_enc<|audio_end|><|im_end|>\n // <|im_start|>assistant\n[language {Name}]? // -// The system prompt is empty (we do not yet thread a caller-supplied -// context through). When `lang_prefix_ids` is non-null its contents -// are appended verbatim after the trailing newline, which is how the -// reference forces a specific output language: the LM continues from -// "" with pure transcript text rather than emitting its own -// "language X" prefix. Callers resolve the prefix via -// encode_language_prefix() below; it stays out of this function so -// build_prompt_tokens can remain a pure token-id assembler. +// System prompt is empty. A non-null `lang_prefix_ids` (resolved via +// encode_language_prefix) is appended after the trailing newline to force an +// output language; kept out of here so this stays a pure token-id assembler. void build_prompt_tokens(const QwenAsrHParams & hp, const ChatTokens & ct, int T_enc, @@ -485,19 +455,13 @@ void build_prompt_tokens(const QwenAsrHParams & hp, } } -} // namespace (close anon temporarily; the two helpers below have - // external linkage — one because the public-header declaration in - // qwen3_asr.h needs a matching definition, the other because the - // function lives just above and wants to name the data table.) - -// BCP-47 → publisher canonical name. The Qwen3-ASR prompt renders the -// publisher's canonical string ("English", "Chinese", ...) rather -// than the BCP-47 code, and the list is frozen per Qwen3-ASR release -// (qwen_asr.inference.utils.SUPPORTED_LANGUAGES in the publisher's -// Python package). Keeping the map here tracks the publisher's list -// directly; if a future Qwen3-ASR variant changes it, this table is -// the update point and the model will error cleanly for languages -// that haven't been added yet. +} // namespace (close anon temporarily for the two external-linkage helpers + // below; encode_language_prefix matches the qwen3_asr.h declaration.) + +// BCP-47 → publisher canonical name ("English", "Chinese", ...), which the +// prompt renders instead of the code. Frozen per release +// (qwen_asr.inference.utils.SUPPORTED_LANGUAGES); this table is the update +// point for a future variant. struct LangNameEntry { const char * bcp47; const char * pub_name; @@ -535,26 +499,12 @@ constexpr LangNameEntry k_qwen3_asr_language_names[] = { {"mk", "Macedonian"}, }; -// Resolve a caller-supplied BCP-47 language code to the token-id -// sequence the chat template expects: the BPE-encoded bytes of -// "language {Name}" followed by the `` special-token id. -// -// We split the prefix around the special token because -// Tokenizer::encode() does not partition special tokens out before -// BPE — it would run a byte-level merge loop over the literal -// "" string and produce wrong ids. The Hugging Face -// tokenizer's "added_tokens" path is what recognizes the tag and -// emits its single id; we replicate that by looking up the id -// directly from the vocab and appending it by hand. The id is -// checkpoint-specific (151704 on Qwen3-ASR 0.6B / 1.7B) but we never -// hardcode it — `tok.find("")` reads the actual vocab. -// -// The dispatcher validates `bcp47` against caps.languages before we -// get here, so an unknown code at this layer is either a converter -// drift (language written into caps that our static map doesn't -// know about) or a future Qwen3-ASR variant we haven't caught up to. -// Either way, surface as UNSUPPORTED_LANGUAGE so the caller sees the -// right error. +// Resolve a caller-supplied BCP-47 code to the token-id sequence the chat +// template expects: BPE("language {Name}") + the `` special id. +// encode() doesn't split special tokens out, so we look up the id +// directly from the vocab (never hardcoded) and append it by hand. The +// dispatcher validates `bcp47` against caps.languages first, so an unknown +// code here means converter/map drift — surface as UNSUPPORTED_LANGUAGE. transcribe_status encode_language_prefix(const transcribe::Tokenizer & tok, const char * bcp47, std::vector & out_ids) @@ -646,26 +596,16 @@ transcribe_status run( return TRANSCRIBE_ERR_INVALID_ARG; } - // Pre-run abort check. Qwen3-ASR is single-shot today; this is - // the single observation point. Stage 2's long-form loop will add - // per-chunk polling for Whisper; the other autoregressive families - // can wire per-step polling when they grow their own long-form - // paths. + // Pre-run abort check (Qwen3-ASR is single-shot, so this is the only + // observation point). if (cc->poll_abort()) { return TRANSCRIBE_ERR_ABORTED; } - // Language hint handling. Null / empty == auto-detect (the LM - // emits its own "language X" prefix which the output - // parser below strips). A non-null code is resolved to the - // token-id sequence for "language {Name}" via the - // tokenizer; we seed the assistant turn with those tokens so the - // LM continues straight into pure transcript text. Dispatcher - // already validated `params->language` against caps.languages, - // so if encode_language_prefix reports "no canonical name" the - // static publisher map and the converter's BCP-47 list have - // drifted — surface that as UNSUPPORTED_LANGUAGE instead of - // silently falling back to auto-detect. + // Language hint. Null/empty == auto-detect (the LM emits its own + // "language X" prefix, stripped by the output parser below). A + // non-null code is resolved to "language {Name}" tokens that seed + // the assistant turn; a resolve failure surfaces as UNSUPPORTED_LANGUAGE. std::vector lang_prefix_ids; const std::vector * lang_prefix_ptr = nullptr; if (params != nullptr && params->language != nullptr && @@ -682,7 +622,7 @@ transcribe_status run( transcribe::debug::init(); - // ----- Mel front-end ------------------------------------------- + // Mel front-end. if (!cm->mel.has_value()) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr run: model has no MelFrontend"); @@ -715,7 +655,7 @@ transcribe_status run( shape, 2, "frontend.mel.norm"); } - // ----- Compute encoder timing + reject unsupported shapes ------ + // Compute encoder timing + reject unsupported shapes. EncoderTiming timing = compute_encoder_timing(mel_n_frames, cm->hparams); if (timing.n_chunks <= 0) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -724,13 +664,10 @@ transcribe_status run( return TRANSCRIBE_ERR_GGUF; } - // ----- Reset per-call compute state ---------------------------- - // Phase timers. Internal diagnostic: the public - // transcribe_timings only exposes mel/encode/decode, but a lot of - // per-run cost sits in graph build, scheduler alloc, host uploads, - // and prefill compute. These locals break that out so we can print - // a full breakdown under TRANSCRIBE_PERF_DEBUG without touching - // the public API. + // Reset per-call compute state. The phase timers below break out + // per-run cost (graph build, sched alloc, uploads, prefill compute) + // that the public transcribe_timings (mel/encode/decode) doesn't + // expose, for the TRANSCRIBE_PERF_DEBUG breakdown. const int64_t t_enc_build_start = ggml_time_us(); int64_t t_enc_build_us = 0; int64_t t_enc_d2h_us = 0; @@ -747,7 +684,7 @@ transcribe_status run( { ggml_init_params ip {}; - ip.mem_size = 16 * 1024 * 1024; // 18 blocks + subsample + head + ip.mem_size = 16 * 1024 * 1024; ip.mem_buffer = nullptr; ip.no_alloc = true; cc->compute_ctx = ggml_init(ip); @@ -758,7 +695,7 @@ transcribe_status run( } } - // ----- Build encoder graph ------------------------------------- + // Build encoder graph. EncoderBuild eb = build_encoder_graph(cc->compute_ctx, cm->weights, cm->hparams, timing, cc->encoder_use_flash); @@ -766,7 +703,7 @@ transcribe_status run( return TRANSCRIBE_ERR_GGUF; } - // ----- Allocate + compute encoder graph ------------------------ + // Allocate + compute encoder graph. if (cc->sched == nullptr) { cc->sched = ggml_backend_sched_new( cm->plan.scheduler_list.data(), nullptr, @@ -815,7 +752,6 @@ transcribe_status run( } } - // Thread count. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); const int64_t t_enc_start = ggml_time_us(); @@ -849,10 +785,9 @@ transcribe_status run( try_dump("enc.ln_post.out", eb.dumps.ln_post_out, "enc.ln_post"); try_dump("enc.proj.out", eb.dumps.proj_out, "enc.proj"); - // Read encoder output to host for use by the LM prefill. The graph - // already dropped the aftercnn pad rows before the encoder blocks - // (see encoder.cpp), so eb.out is exactly [d_enc, T_enc] — matching - // the reference's `padded_embed[padded_mask_after_cnn]`-selected + // Read encoder output to host for the LM prefill. The graph already + // dropped the aftercnn pad rows (see encoder.cpp), so eb.out is exactly + // [d_enc, T_enc] — the reference's `padded_embed[padded_mask_after_cnn]` // shape. const int d_enc = static_cast(eb.out->ne[0]); const int T_enc = static_cast(eb.out->ne[1]); @@ -863,14 +798,12 @@ transcribe_status run( cc->enc_host.size() * sizeof(float)); t_enc_d2h_us = ggml_time_us() - t_d2h_start; - // ----- Decode phase begins ----- - // t_dec_start covers the whole decode phase (prompt + KV init + - // prefill build/compute + step loop). Prefill is part of "decode" as - // users understand it. + // Decode phase begins. t_dec_start covers prompt + KV init + prefill + // build/compute + step loop (prefill is part of "decode" to users). const int64_t t_dec_start = ggml_time_us(); const int64_t t_prefill_build_start = t_dec_start; - // ----- Prompt construction ----- + // Prompt construction. std::vector prompt_ids; std::vector audio_positions; build_prompt_tokens(cm->hparams, cm->chat_tokens, T_enc, @@ -879,13 +812,10 @@ transcribe_status run( const int prefix_len = audio_positions.empty() ? 0 : static_cast(audio_positions.front()); const int suffix_len = T_prompt - prefix_len - T_enc; - (void)audio_positions; // no longer fed to the graph + (void)audio_positions; - // ----- Input-length gate (see docs/input-limits.md) ----- - // audio tokens + prompt + generation must fit the decoder context - // window (optionally lowered by the caller's n_ctx). The audio-token - // count is fixed by the input length, so reject an over-length clip - // here, before prefill/decode, instead of walling at a fixed size. + // Input-length gate: audio + prompt + generation must fit the decoder + // context window. Reject an over-length clip here, before prefill/decode. const int ceiling = qwen3_context_ceiling(cc->n_ctx, cm->hparams); if (T_prompt + k_max_new > ceiling) { transcribe::log_msg( @@ -898,13 +828,10 @@ transcribe_status run( return TRANSCRIBE_ERR_INPUT_TOO_LONG; } - // ----- KV cache init (grow-to-fit, clamped to the context ceiling) ----- - // Size to hold the prompt plus the generation budget, rounded up to a - // 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. For - // typical short clips this is the same 1024/2048 width as before, so the - // decode is unchanged — only previously-walled long clips now fit. + // 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 + // 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) want_n_ctx *= 2; if (want_n_ctx > ceiling) want_n_ctx = ceiling; @@ -938,11 +865,9 @@ transcribe_status run( cc->kv_cache.head = 0; } - // ----- Prefill graph ----- - // slice_last: when false, the last block's FFN + final norm run - // on every position (needed for debug dump parity). When true, we - // slice to just the final position before the last FFN — same - // optimization llama.cpp's inp_out_ids gives, worth ~25 ms here. + // Prefill graph. slice_last false: last block's FFN + final norm run on + // every position (needed for dump parity). true: slice to just the final + // position before the last FFN (llama.cpp's inp_out_ids trick, ~25 ms). const bool dumps_on = transcribe::debug::enabled(); const bool slice_last = !dumps_on; PrefillBuild pb = build_prefill_graph( @@ -977,9 +902,8 @@ transcribe_status run( } { - // Causal mask in F16 (matches pb.mask_in's type). Row r, col c: - // +0 if c <= r, else -inf. Row-major (ne[0]=T_prompt fastest). - // F16 direct upload avoids a per-layer ggml_cast in the graph. + // Causal mask in F16 (matches pb.mask_in): row r col c is 0 if c <= r, + // else -inf, row-major. F16 upload avoids a per-layer ggml_cast. const ggml_fp16_t mask_zero = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t mask_neg_inf = ggml_fp32_to_fp16(-INFINITY); std::vector mask( @@ -1022,7 +946,7 @@ transcribe_status run( try_dump("dec.out_before_head", pb.dumps.out_before_head, "dec.out_before_head"); try_dump("dec.logits_raw", pb.dumps.logits_raw, "dec.logits_raw"); - // ----- Read prefill logits + first argmax ----- + // Read prefill logits + first argmax. const int64_t t_prefill_logits_start = ggml_time_us(); const int vocab = cm->hparams.dec_vocab_size; std::vector logits(vocab); @@ -1043,18 +967,15 @@ transcribe_status run( generated_ids.push_back(next_tok); t_prefill_logits_us = ggml_time_us() - t_prefill_logits_start; - // ----- Step loop ----- + // Step loop. const int32_t eos_id = cm->hparams.eos_token_id; - const int32_t max_new = k_max_new; // matches reference dumper default + const int32_t max_new = k_max_new; int cur_past = T_prompt; - // ---------- Build step graph ONCE and reuse every step ---------- - // Static shape sized for the actual workload: T_prompt audio/prompt - // tokens already written + up to max_new generated tokens. - // Metal's flash-attn kernels dispatch much faster (~30% on M4 Max) - // when the K/V ne[1] is a power of 2. Round up accordingly, with a - // floor of 1024 since smaller values just move us into the "slow - // misaligned" branch without saving meaningful bandwidth. + // Build the step graph ONCE and reuse every step, sized for the actual + // workload (T_prompt written + up to max_new generated). Metal's flash-attn + // kernels dispatch ~30% faster (M4 Max) when K/V ne[1] is a power of 2, so + // round up; floor of 1024 (smaller just hits the slow-misaligned branch). int max_n_kv = 1024; while (max_n_kv < T_prompt + max_new) max_n_kv *= 2; if (max_n_kv > cc->kv_cache.n_ctx) max_n_kv = cc->kv_cache.n_ctx; @@ -1091,14 +1012,13 @@ transcribe_status run( } const int64_t t_step_build_once_us = ggml_time_us() - t_step_build_start; - // Mask buffer: reused host-side across steps. Starts all -inf; - // per step we zero positions [0, cur_past] (attend) and leave - // [cur_past+1, max_n_kv) as -inf (ignore). + // Mask buffer reused host-side across steps. Starts all -inf; per step we + // zero positions [0, cur_past] (attend) and leave the rest -inf. const ggml_fp16_t mask_zero = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t mask_neg_inf = ggml_fp32_to_fp16(-INFINITY); std::vector step_mask(max_n_kv, mask_neg_inf); - // Per-step timers. Now much smaller — no per-step graph work. + // Per-step timers. int64_t t_step_set_us = 0; int64_t t_step_comp_us = 0; int64_t t_step_get_us = 0; @@ -1158,10 +1078,8 @@ transcribe_status run( t_step_loop_us = ggml_time_us() - t_step_loop_start; n_steps = static_cast(generated_ids.size()) - 1; - // The decode stopped at EOS (complete) or at the generation budget / - // context width (truncated). Surface the latter via - // transcribe_was_truncated() and a WARN rather than returning a - // silently shortened transcript. See docs/input-limits.md. + // Decode stopped at EOS (complete) or the generation budget / context width + // (truncated). Surface the latter via transcribe_was_truncated() + WARN. if (next_tok != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -1172,10 +1090,9 @@ transcribe_status run( static_cast(generated_ids.size())); } - // Map the granular diagnostic counters to the shape the debug - // print expects. With graph reuse all per-step overhead collapses - // to tensor_set; build/alloc/ctx_reset are amortized in the - // t_step_build_once_us bucket printed separately. + // Map granular counters to the debug-print shape. With graph reuse all + // per-step overhead collapses to tensor_set; build/alloc/ctx_reset are + // amortized in t_step_build_once_us. int64_t t_step_ctx_us = 0; int64_t t_step_build_us = t_step_build_once_us; int64_t t_step_alloc_us = 0; @@ -1185,30 +1102,20 @@ transcribe_status run( generated_ids.pop_back(); } - // Decode generated ids to text. Tokenizer::decode handles the - // "gpt2" byte-level inversion natively — no more per-family - // byte-to-unicode walker. + // Decode generated ids to text (Tokenizer::decode handles the "gpt2" + // byte-level inversion natively). std::string raw_text = cm->tok.decode( generated_ids.data(), static_cast(generated_ids.size())); - // Parse Qwen3-ASR output format: - // auto-detect: "language Xactual_text" (strip prefix) - // forced: "actual_text" (no prefix - - // we seeded the assistant turn with "language - // X" already, so the LM's generated - // continuation is pure transcript.) - // - // We unconditionally try the split — if the prefix is absent (the - // forced case, or a model that just didn't emit it this run), the - // find() returns npos and transcript_text stays as raw_text. + // Parse Qwen3-ASR output: auto-detect emits "language Xtext" + // (strip prefix); forced emits "text" (we already seeded the prefix). The + // split is unconditional — absent prefix leaves transcript_text = raw_text. std::string transcript_text = raw_text; if (auto sep = raw_text.find(""); sep != std::string::npos) { - // Auto-detect path: the prefix carries the language name the - // model picked. Surface it as detected_language (reverse-map - // the human-readable name to its BCP-47 code via the same - // table encode_language_prefix uses), but only when the caller - // did NOT supply a hint — the public field's contract is "what - // the model told us," not "what we told the model." + // Auto-detect path: surface the model-picked language name as + // detected_language (reverse-mapped to BCP-47), but only when the + // caller did NOT supply a hint (the field reports what the model told + // us, not what we told it). if (lang_prefix_ptr == nullptr) { constexpr const char k_prefix[] = "language "; const size_t name_start = raw_text.find(k_prefix); @@ -1231,15 +1138,11 @@ transcribe_status run( transcript_text = raw_text.substr(sep + std::strlen("")); } - // Write full_text + a single segment (no timestamps; Qwen3-ASR's - // ASR head emits TIMESTAMPS_NONE per the family capability). + // Write full_text + a single segment (no timestamps; TIMESTAMPS_NONE). cc->t_decode_us = ggml_time_us() - t_dec_start; - // Optional perf breakdown. Public transcribe_timings only has - // mel/encode/decode — this dumps the finer split we need to - // reason about bench wall-vs-phase gaps without changing the - // public API or the bench JSON schema. Gated on env var so it - // doesn't spam normal runs. + // Optional perf breakdown (finer split than the public mel/encode/decode + // timings), gated on env var. if (const char * e = std::getenv("TRANSCRIBE_PERF_DEBUG"); e != nullptr && *e != '\0' && *e != '0') { @@ -1295,9 +1198,8 @@ transcribe_status run( / static_cast(cm->hparams.fe_sample_rate); cc->segments.push_back(std::move(seg)); - // The partial transcript is fully populated above; a truncated decode - // returns the hard OUTPUT_TRUNCATED status (the result stays readable, - // like an aborted run). See docs/input-limits.md. + // A truncated decode returns OUTPUT_TRUNCATED; the partial transcript above + // stays readable (like an aborted run). return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } @@ -1306,12 +1208,10 @@ transcribe_status run( // Offline batched decode (transcribe_run_batch) // =========================================================================== // -// Strategy (see docs/batching-autoregressive-plan.md): prefill each utterance -// serially into its OWN slab of a batched KV cache — byte-identical to the -// single-shot prefill, so it inherits correctness — then batch only the -// expensive autoregressive step loop. The encoder + prefill stay per-utterance -// (low risk); all new batch-axis math lives in the batched step graph and is -// gated by per-row text/logits parity vs single-shot. +// Prefill each utterance serially into its OWN slab of a batched KV cache +// (byte-identical to single-shot prefill, so it inherits correctness), then +// batch only the autoregressive step loop. Encoder + prefill stay +// per-utterance; the batch-axis math lives in the batched step graph. namespace { @@ -1347,7 +1247,7 @@ transcribe_status encode_all_batched( int64_t & mel_us, int64_t & enc_us) { const int n_mels = cm->hparams.enc_num_mel_bins; - // ---- mel (parallel across utterances) ---- + // Mel (parallel across utterances). std::vector> mels(n); std::vector mel_nf(n, 0); int n_threads = cc->n_threads; @@ -1368,7 +1268,7 @@ transcribe_status encode_all_batched( }); mel_us += ggml_time_us() - t_mel0; - // ---- per-utterance timing + packing geometry ---- + // Per-utterance timing + packing geometry. std::vector timings(n); int n_chunks_max = 1; bool any = false; @@ -1382,7 +1282,7 @@ transcribe_status encode_all_batched( } if (!any) return TRANSCRIBE_OK; // caller emits per-row errors - // ---- build batched encoder graph ---- + // Build batched encoder graph. if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) return st; EncoderBuildBatched eb = build_encoder_graph_batched( @@ -1502,7 +1402,7 @@ transcribe_status encode_all_batched( // KV into its slab and returning the first generated token per utterance. The // caller must have allocated cc->kv_cache_batch with n_ctx >= max T_prompt and // computed prompt_ids[b] / T_prompt[b] / prefix_len. Collapses B per-utterance -// prefill graph-builds + sched-allocs into one (the L40S floor-lift). +// prefill graph-builds + sched-allocs into one. transcribe_status prefill_all_batched( QwenAsrSession * cc, QwenAsrModel * cm, const std::vector & valid, @@ -1545,10 +1445,9 @@ transcribe_status prefill_all_batched( ggml_backend_tensor_set(pb.input_ids_in, ids.data(), 0, ids.size() * sizeof(int32_t)); } - // Audio injection by elementwise blend: audio_dense [d_enc, T_prompt_max, n] - // holds each utterance's enc_out embeds scattered into their prompt - // positions (zero elsewhere), keep [T_prompt_max, n] is 0 there and 1 - // elsewhere. x = x*keep + audio_dense. (d_enc == dec_hidden, enforced.) + // Audio injection by elementwise blend (see decoder.h): audio_dense holds + // each utterance's enc_out embeds scattered into their prompt positions, + // keep is 0 there and 1 elsewhere. (d_enc == dec_hidden, enforced.) { std::vector audio_dense( static_cast(d_enc) * T_prompt_max * n, 0.0f); @@ -1708,8 +1607,7 @@ transcribe_status run_batch( if (!cm->mel.has_value()) return TRANSCRIBE_ERR_INVALID_ARG; // Batched decode requires the flash-attention step path and dump-free - // operation (the per-utterance graphs reuse one compute_ctx). Fall back - // to the serial loop otherwise — same results, no throughput gain. + // operation. Fall back to the serial loop otherwise (same results). if (!cc->decoder_use_flash || transcribe::debug::enabled() || n == 1) { return run_batch_serial(cc, pcm, n_samples, n, params); } @@ -1723,13 +1621,12 @@ transcribe_status run_batch( params->language[0] != '\0') { if (encode_language_prefix(cm->tok, params->language, lang_prefix_ids) != TRANSCRIBE_OK) { - // Bad language hint is a whole-batch config error; mirror run(). return TRANSCRIBE_ERR_UNSUPPORTED_LANGUAGE; } lang_prefix_ptr = &lang_prefix_ids; } - // ---- Pass 1: per-utterance encoder + prefill into KV slabs ---- + // Pass 1: per-utterance encoder + prefill into KV slabs. std::vector> generated(n); std::vector T_prompt(n, 0); std::vector next_tok(n, 0); @@ -1754,10 +1651,8 @@ transcribe_status run_batch( const int max_new = 256; int max_T_prompt = 0; int prefix_len = 0; - // Per-utterance terminal status for rows we reject before decode. Defaults - // to INVALID_ARG (the encode pass already invalidated malformed rows); - // over-length rows below are upgraded to INPUT_TOO_LONG so run_batch - // enforces the same input-length contract as the single-shot run(). + // Per-utterance terminal status for rejected rows. Defaults to INVALID_ARG; + // over-length rows below are upgraded to INPUT_TOO_LONG. const int ceiling = qwen3_context_ceiling(cc->n_ctx, cm->hparams); std::vector fail_status(n, TRANSCRIBE_ERR_INVALID_ARG); std::vector> prompt_ids(n); @@ -1768,9 +1663,7 @@ transcribe_status run_batch( lang_prefix_ptr, prompt_ids[b], ap); T_prompt[b] = static_cast(prompt_ids[b].size()); prefix_len = ap.empty() ? 0 : static_cast(ap.front()); - // Same gate as single-shot run(): audio + prompt + generation must fit - // the context ceiling. Reject the over-length utterance (the rest of - // the batch still runs); see docs/input-limits.md. + // Same gate as single-shot run(); the rest of the batch still runs. if (T_prompt[b] + max_new > ceiling) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "qwen3_asr run_batch: utterance %d input too long — %d audio + " @@ -1794,9 +1687,8 @@ transcribe_status run_batch( } int max_n_kv = 1024; while (max_n_kv < max_T_prompt + max_new) max_n_kv *= 2; - // Honor the session context cap: the pow2 rounding above can overshoot the - // ceiling, so clamp (the per-utterance gate guarantees every valid row's - // T_prompt + max_new <= ceiling, so the cache still fits all rows). + // Clamp the pow2 round-up to the ceiling (the per-utterance gate guarantees + // every valid row still fits). if (max_n_kv > ceiling) max_n_kv = ceiling; // Allocate / reuse the batched KV cache (n_ctx == max_n_kv, n slabs). @@ -1839,7 +1731,7 @@ transcribe_status run_batch( } const int64_t prefill_pass_us = ggml_time_us() - t_prefpass0; - // ---- Pass 2: batched step loop (shared causal_lm driver) ---- + // Pass 2: batched step loop (shared causal_lm driver). const int32_t eos_id = cm->hparams.eos_token_id; if (const transcribe_status st = reset_compute_ctx(cc, 16); st != TRANSCRIBE_OK) @@ -1875,7 +1767,7 @@ transcribe_status run_batch( const int64_t step_us = step_stats.step_us; const int n_steps = step_stats.n_steps; - // ---- Capture per-utterance results ---- + // Capture per-utterance results. const int valid_count = std::max(1, static_cast( std::count(valid.begin(), valid.end(), char(1)))); for (int b = 0; b < n; ++b) { @@ -1887,9 +1779,7 @@ transcribe_status run_batch( } transcribe_session::ResultSet rs = finalize_utterance(cm, generated[b], lang_prefix_ptr, n_samples[b]); - // Per-utterance truncation parity with the single-shot path: a row that - // hit the generation budget / context before eos reports - // TRANSCRIBE_ERR_OUTPUT_TRUNCATED (partial transcript retained). + // Per-utterance truncation parity with the single-shot path. if (b < static_cast(truncated.size()) && truncated[b]) { cc->was_truncated = true; rs.status = TRANSCRIBE_ERR_OUTPUT_TRUNCATED; diff --git a/src/arch/qwen3_asr/qwen3_asr.h b/src/arch/qwen3_asr/qwen3_asr.h index 5d770ad1..a40974b1 100644 --- a/src/arch/qwen3_asr/qwen3_asr.h +++ b/src/arch/qwen3_asr/qwen3_asr.h @@ -1,13 +1,8 @@ // arch/qwen3_asr/qwen3_asr.h - Qwen3-ASR model and context types. // -// INTERNAL to src/arch/qwen3_asr/. Defines the concrete classes that -// derive from transcribe_model / transcribe_session for the Qwen3-ASR -// family (audio-LLM: audio encoder + Qwen3 causal LM with audio-token -// injection). -// -// First port targets non-streaming ASR transcription. The sibling -// forced-aligner variant (Qwen3-ForcedAligner-0.6B) and the streaming -// stream_transcribe path are tracked separately. +// INTERNAL to src/arch/qwen3_asr/. Concrete transcribe_model / +// transcribe_session subclasses for the Qwen3-ASR family (audio-LLM: audio +// encoder + Qwen3 causal LM with audio-token injection). #pragma once @@ -40,12 +35,10 @@ namespace transcribe::qwen3_asr { void apply_family_invariants(transcribe_model & model); -// Encode "language {Name}" for the given BCP-47 code. The -// output is the token-id sequence the Qwen3-ASR chat template seeds -// the assistant turn with when a caller forces a language hint. -// Declared here (rather than hiding in the model.cpp anonymous -// namespace) so the BPE parity test can verify the full prefix -// matches the HF reference including the special token. +// Encode "language {Name}" for the given BCP-47 code: the token-id +// sequence the chat template seeds the assistant turn with on a forced hint. +// Declared here (not in model.cpp's anon namespace) so the BPE parity test can +// verify the full prefix against the HF reference. // // Returns: // TRANSCRIBE_OK on success, out_ids populated. @@ -66,16 +59,9 @@ transcribe_status encode_language_prefix(const transcribe::Tokenizer & tok, // Model / Context // --------------------------------------------------------------------------- -// Qwen3 chat-template special-token ids resolved through the loaded -// tokenizer at load time. The prompt renderer is narrow (no Jinja) -// but it still tokenizes to exact ids; resolving at load lets a -// future checkpoint that reorders the vocab fail loudly instead of -// silently corrupting the prompt. -// -// All fields are the *resolved* ids; the reference ids on the 0.6B -// and 1.7B vocab are stable (im_start=151644, im_end=151645, -// newline=198, system=8948, user=872, assistant=77091) but we do -// not hardcode them at runtime. +// Qwen3 chat-template special-token ids, resolved through the loaded tokenizer +// at load time so a future vocab reorder fails loudly. Reference ids on the +// 0.6B/1.7B vocab are stable but never hardcoded at runtime. struct ChatTokens { int32_t im_start = -1; int32_t im_end = -1; @@ -97,14 +83,10 @@ struct QwenAsrModel final : public transcribe_model { std::optional mel; - // Jinja chat template. The prompt is rebuilt at run time with the - // caller's language / context; we only store the template string - // from the GGUF KV. Empty string means "no template" (the first - // port will refuse to run without one). + // Jinja chat template string from the GGUF KV (empty == none). std::string chat_template; - // Resolved chat-template token ids (filled at load time, see - // resolve_chat_tokens in model.cpp). Used by build_prompt_tokens. + // Resolved chat-template token ids (see resolve_chat_tokens in model.cpp). ChatTokens chat_tokens; QwenAsrModel() = default; diff --git a/src/arch/sensevoice/capabilities.cpp b/src/arch/sensevoice/capabilities.cpp index d4060628..316eb02e 100644 --- a/src/arch/sensevoice/capabilities.cpp +++ b/src/arch/sensevoice/capabilities.cpp @@ -22,19 +22,13 @@ void apply_family_invariants(transcribe_model & model) { // The non-AR CTC head has no segment- or word-timestamp head in // the published runtime. forced-CTC alignment IS available via // ctc_forced_align in the upstream demo, but the C++ runtime does - // not expose it in v1; mark NONE and revisit in a follow-up if a - // user asks for word timing. + // not expose it; mark NONE. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_NONE; // Feature bits: cancellation is wired at the run level. SenseVoice // exposes a runtime ITN toggle via the textnorm prefix embedding // (`woitn` / `withitn`); the generic transcribe_run_params::itn enum // routes here. No PNC runtime toggle. - // - // TODO(family doc): on en/zh/yue/ja/ko, observe whether `withitn` - // also bundles punctuation/casing changes alongside number/date - // normalization and capture in the family doc. API shape doesn't - // change either way. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_CANCELLATION, true); transcribe::set_feature(&model, TRANSCRIBE_FEATURE_ITN, true); } diff --git a/src/arch/sensevoice/encoder.cpp b/src/arch/sensevoice/encoder.cpp index b7962450..2fe58946 100644 --- a/src/arch/sensevoice/encoder.cpp +++ b/src/arch/sensevoice/encoder.cpp @@ -1,26 +1,11 @@ // arch/sensevoice/encoder.cpp - SenseVoice SAN-M encoder graph builder. // -// Forward shape (single-utterance, batch=1, channel-innermost ggml ne): -// -// frontend.in [d_input=560, T_lfr] -// -> prefix prepend (lid + event_emo + textnorm via ggml_get_rows) -// [d_input, 4 + T_lfr] = enc.input.with_prefix -// -> scale by sqrt(d_model=512) -// -> add sinusoidal PE (depth = current width = d_input = 560, -// 1-based positions) -// [d_input, 4 + T_lfr] = enc.embed.out -// -> encoders0[0] (SAN-M block, projection 560 -> 512, no attn-residual) -// [d_model, 4 + T_lfr] = enc.encoders0.0.out -// -> encoders[0..n-2] x49 (SAN-M block, in==out==d_model) -// [d_model, 4 + T_lfr] = enc.encoders.{0,24,48}.out -// -> after_norm (LayerNorm eps=1e-12) -// [d_model, 4 + T_lfr] = enc.after_norm.out -// -> tp_encoders[0..tp-1] x20 (SAN-M block) -// [d_model, 4 + T_lfr] = enc.tp_encoders.{0,10,19}.out -// -> tp_norm (LayerNorm eps=1e-12) -// [d_model, 4 + T_lfr] = enc.tp_norm.out -// -> ctc.head linear [d_model, vocab] + bias [vocab] -// [vocab, 4 + T_lfr] = ctc.logits.raw +// Forward (single-utterance, channel-innermost ggml ne; T = 4 + T_lfr after +// the prefix prepend): +// frontend.in [d_input=560, T_lfr] -> prefix prepend (lid + event_emo + +// textnorm) -> scale by sqrt(d_model=512) + sinusoidal PE -> +// encoders0[0] (SAN-M projection 560->512) -> encoders x49 -> +// after_norm -> tp_encoders x20 -> tp_norm -> ctc.head [vocab, T]. // // SAN-M block topology lives in src/sanm/sanm.h (shared with funasr_nano). // This file owns the surrounding shape — prefix prepend, sinusoidal PE @@ -69,9 +54,8 @@ sanm::SanmBlockView to_sanm_view(const SenseVoiceBlock & b) { return v; } -// Mark a tensor as live for the post-compute dump pass and stash the -// pointer in the dumps slot. The mark is a no-op when -// TRANSCRIBE_DUMP_DIR is unset. +// Name a tensor, mark it for the post-compute dump pass (no-op when +// TRANSCRIBE_DUMP_DIR is unset), and stash the pointer in the dumps slot. void mark_dump(ggml_tensor *& slot, ggml_tensor * t, const char * name) { named(t, name); transcribe::debug::mark_tensor_for_dump(t); diff --git a/src/arch/sensevoice/encoder.h b/src/arch/sensevoice/encoder.h index c9c34385..0cd1dc35 100644 --- a/src/arch/sensevoice/encoder.h +++ b/src/arch/sensevoice/encoder.h @@ -1,9 +1,8 @@ // arch/sensevoice/encoder.h - SenseVoice SAN-M encoder graph builder. // -// Stage 4 brings up the encoder graph here. The .cpp file builds a -// ggml_cgraph that mirrors SenseVoiceEncoderSmall.forward (SAN-M -// blocks with FSMN memory branch, two-tier depth, sinusoidal PE, -// pre-encoder prefix-embedding prepend). +// The .cpp file builds a ggml_cgraph that mirrors +// SenseVoiceEncoderSmall.forward (SAN-M blocks with FSMN memory branch, +// two-tier depth, sinusoidal PE, pre-encoder prefix-embedding prepend). // // This header is INTERNAL to src/arch/sensevoice/. The C ABI sees // only SenseVoice::run, which builds and runs this graph. @@ -21,9 +20,8 @@ namespace transcribe::sensevoice { struct SenseVoiceHParams; struct SenseVoiceWeights; -// Named intermediate dump points for the validate.py harness. -// Mirrors dump_coverage.json. Each field is a borrowed pointer into -// the compute_ctx; nullptr means the sub-stage is not yet wired. +// Named intermediate dump points. Each field is a borrowed pointer into +// the compute_ctx; nullptr means the sub-stage is not wired. struct EncoderDumps { // Frontend output (LFR + CMVN). ne=[d_input, T_lfr, 1, 1]. ggml_tensor * frontend_out = nullptr; diff --git a/src/arch/sensevoice/model.cpp b/src/arch/sensevoice/model.cpp index eba420d3..03a6eec0 100644 --- a/src/arch/sensevoice/model.cpp +++ b/src/arch/sensevoice/model.cpp @@ -454,7 +454,6 @@ transcribe_status run( const int T_full = T_lfr + 4; // 4 prepended prefix tokens - // ---------- Reset per-call compute state -------------------------- if (cc->compute_ctx != nullptr) { ggml_free(cc->compute_ctx); cc->compute_ctx = nullptr; @@ -520,8 +519,9 @@ transcribe_status run( // ITN slot. Generic transcribe_run_params::itn routes here. DEFAULT maps // to the shipped behavior (use_itn=false; matches the family's // `itn=False` Python default). OFF / ON override explicitly. The - // dispatcher's advisory WARN only fires when supports_itn == false; - // SenseVoice advertises supports_itn = true, so no WARN fires here. + // dispatcher's advisory WARN only fires when transcribe_model_supports( + // model, TRANSCRIBE_FEATURE_ITN) is false; SenseVoice sets + // TRANSCRIBE_FEATURE_ITN so the probe returns true and no WARN fires here. bool use_itn = false; if (params != nullptr) { switch (params->itn) { diff --git a/src/arch/sensevoice/weights.cpp b/src/arch/sensevoice/weights.cpp index a1123570..aa30ae21 100644 --- a/src/arch/sensevoice/weights.cpp +++ b/src/arch/sensevoice/weights.cpp @@ -1,12 +1,5 @@ // arch/sensevoice/weights.cpp - read_sensevoice_hparams + // build_sensevoice_weights. -// -// Pattern mirrors src/arch/parakeet/weights.cpp: -// - read every required hparam from KV explicitly, -// - build the tensor catalog as get_tensor() calls with explicit -// expected shapes, -// - missing tensor or shape mismatch → TRANSCRIBE_ERR_GGUF with a -// log naming the offending tensor. #include "weights.h" @@ -207,9 +200,8 @@ transcribe_status load_block(ggml_context * ctx_meta, std::snprintf(buf, sizeof(buf), "%s.%s", prefix, suffix); return buf; }; - // Reserve a second buffer for callers that need two names live in - // the same expression. We dispatch through copies into local - // std::strings to dodge buffer reuse. + // Names are built as local std::strings so two can be live in one + // expression without the single `buf` aliasing. std::string name_norm_attn_w = std::string(prefix) + ".norm_attn.weight"; std::string name_norm_attn_b = std::string(prefix) + ".norm_attn.bias"; std::string name_qkv_w = std::string(prefix) + ".attn.qkv.weight"; @@ -243,7 +235,7 @@ transcribe_status load_block(ggml_context * ctx_meta, GET_F32(b.ffn_fc1_b, name_fc1_b.c_str(), d_ff); GET_LIN(b.ffn_fc2_w, name_fc2_w.c_str(), d_ff, d_model); GET_F32(b.ffn_fc2_b, name_fc2_b.c_str(), d_model); - (void)namef; // currently unused; left for future per-block helpers. + (void)namef; return TRANSCRIBE_OK; } diff --git a/src/arch/voxtral/decoder.cpp b/src/arch/voxtral/decoder.cpp index 2193eb8d..d705d8ea 100644 --- a/src/arch/voxtral/decoder.cpp +++ b/src/arch/voxtral/decoder.cpp @@ -1,11 +1,10 @@ // arch/voxtral/decoder.cpp - Voxtral LM prefill/step graph builders. // -// Reference: VoxtralForConditionalGeneration's inner language_model -// (LlamaForCausalLM). The per-block math (pre-LN RMSNorm, GQA, NEOX -// RoPE, SwiGLU on packed gate_up) is the shared causal_lm module, called -// 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 head. +// Reference: VoxtralForConditionalGeneration's inner LlamaForCausalLM. The +// per-block math (pre-LN RMSNorm, GQA, NEOX RoPE, SwiGLU on packed gate_up) 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. #include "decoder.h" @@ -296,8 +295,8 @@ StepBuild build_step_graph(ggml_context * ctx, ggml_tensor * logits = ggml_mul_mat(ctx, weights.dec_embed.output_w, x); logits = ggml_reshape_1d(ctx, logits, vocab); ggml_set_name(logits, "step.logits"); - // Kept as a graph output so the mid-generation dec.logits_raw.gen - // dump can read it back after a step compute (validate.py coverage). + // Kept as a graph output so the mid-generation dec.logits_raw.gen dump + // can read it back after a step compute. ggml_set_output(logits); sb.logits = logits; @@ -352,10 +351,8 @@ PrefillBuildBatched build_prefill_graph_batched( (void) T_audio_max; // audio_dense is prompt-sized; T_audio_max no longer used pb.input_ids_in = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, T_prompt_max, B); ggml_set_input(pb.input_ids_in); - // Audio injection is an elementwise blend over the flat token axis (no - // set_rows): a k-quant token_embd get_rows is unsupported on CUDA and runs - // on CPU; a set_rows consuming that CPU tensor straddles the CPU/CUDA split - // and faults. x*keep_mask + audio_dense crosses the split cleanly. + // Audio injection is an elementwise blend (see header): x*keep_mask + + // audio_dense, which crosses the CPU/CUDA split cleanly unlike a set_rows. pb.audio_dense_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, static_cast(T_prompt_max) * B); ggml_set_input(pb.audio_dense_in); @@ -378,14 +375,12 @@ PrefillBuildBatched build_prefill_graph_batched( } pb.graph = gf; - // Token-embed all (right-padded) prompt rows, then scatter the audio - // embeddings over the audio_token_id placeholder positions. + // Token-embed all (right-padded) prompt rows, then blend the audio embeds in + // elementwise: x = x*keep_mask + audio_dense (keep_mask broadcasts over + // hidden). Reshape to 3D for the batched block stack. ggml_tensor * ids_flat = ggml_reshape_1d(ctx, pb.input_ids_in, static_cast(T_prompt_max) * B); - // x is 2D [hidden, T_prompt_max*B] (contiguous). Blend the audio embeds in - // elementwise: x = x*keep_mask + audio_dense. keep_mask broadcasts over - // hidden (ne[0]). Then reshape to 3D for the batched block stack. ggml_tensor * x = ggml_get_rows(ctx, weights.dec_embed.token_w, ids_flat); x = ggml_add(ctx, ggml_mul(ctx, x, pb.keep_mask_in), pb.audio_dense_in); x = ggml_reshape_3d(ctx, x, hidden, T_prompt_max, B); diff --git a/src/arch/voxtral/decoder.h b/src/arch/voxtral/decoder.h index 68e4918c..b5cbac07 100644 --- a/src/arch/voxtral/decoder.h +++ b/src/arch/voxtral/decoder.h @@ -8,13 +8,13 @@ // - UNTIED lm_head (dec.output.weight, separate from token_embd) // // The per-block math is the shared causal_lm module; this file owns graph -// allocation, prompt + audio injection (3-way concat of -// prefix | proj.out | suffix), tensor naming / dump parity, the untied -// head, and the driver-facing build structs. +// allocation, audio injection, tensor naming / dump parity, the untied head, and +// the driver-facing build structs. // -// Audio injection: the prompt has `audio_token_id` placeholders at a -// single contiguous run [prefix_len, prefix_len + T_enc); the projector -// output (T_enc audio embeddings in dec_hidden space) is spliced there. +// Audio injection: the prompt has `audio_token_id` placeholders at a single +// contiguous run [prefix_len, prefix_len + T_enc); the projector output (T_enc +// audio embeddings) is spliced there via a 3-way concat (prefix | proj.out | +// suffix). #pragma once @@ -92,14 +92,12 @@ StepBuild build_step_graph(ggml_context * ctx, // --------------------------------------------------------------------------- // Batched prefill / step (offline transcribe_run_batch fast path) // -// B utterances decoded in lockstep against a batched KV cache -// (causal_lm::kv_init_batched, n_batch=B). Ragged prompts are right-padded -// to T_prompt_max; the shared causal mask works because pads land after -// each utterance's real tokens, and per-row last_idx selects each real -// final position. Audio embeddings are scattered into the token-embedding -// sequence via ggml_set_rows at the audio_token_id positions. Both require -// use_flash (the causal_lm batched blocks have no manual GQA path). Mirrors -// arch/canary_qwen, with Voxtral's UNTIED lm_head (dec.output.weight). +// B utterances decoded in lockstep against a batched KV cache (n_batch=B). +// Ragged prompts are right-padded to T_prompt_max; the shared causal mask works +// because pads land after each utterance's real tokens, and per-row last_idx +// selects each real final position. Audio embeddings are blended in at the +// audio_token_id positions. Both require use_flash (the batched blocks have no +// manual GQA path). UNTIED lm_head (dec.output.weight). // --------------------------------------------------------------------------- struct PrefillBuildBatched { diff --git a/src/arch/voxtral/encoder.cpp b/src/arch/voxtral/encoder.cpp index 502816d6..b5af82d2 100644 --- a/src/arch/voxtral/encoder.cpp +++ b/src/arch/voxtral/encoder.cpp @@ -1,7 +1,7 @@ // arch/voxtral/encoder.cpp - Voxtral audio encoder + projector graph. // // Reference: VoxtralEncoder (= Whisper-large-v3 encoder) + -// VoxtralMultiModalProjector in modeling_voxtral.py. +// VoxtralMultiModalProjector. // // mel [n_mels=128, T=3000] // -> transpose to [T, n_mels] for ggml_conv_1d diff --git a/src/arch/voxtral/encoder.h b/src/arch/voxtral/encoder.h index a7a2aadc..024b0fa2 100644 --- a/src/arch/voxtral/encoder.h +++ b/src/arch/voxtral/encoder.h @@ -49,14 +49,12 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, int n_mel_frames, bool use_flash); -// Batched encoder+projector: process `n_chunks` equal-length 30 s chunks -// at once (batch on ne[2]). mel_in is [n_mels, n_mel_frames, n_chunks]; -// out (proj.out) is [dec_hidden, n_audio_tokens, n_chunks]. Each chunk is -// an independent forward pass sharing weights, so chunk b is numerically -// identical to the single-chunk build_encoder_graph on chunk b (the batch -// parity gate relies on this). Requires use_flash (no batched manual-GQA -// path); the caller falls back to per-chunk when flash is off. No dump -// hooks — the batch fast path never runs under debug dumps. +// Batched encoder+projector: process `n_chunks` equal-length 30 s chunks at once +// (batch on ne[2]). mel_in is [n_mels, n_mel_frames, n_chunks]; out (proj.out) is +// [dec_hidden, n_audio_tokens, n_chunks]. Each chunk is an independent forward +// pass sharing weights, so chunk b is numerically identical to the single-chunk +// build_encoder_graph. Requires use_flash (no batched manual-GQA path); the +// caller falls back to per-chunk when flash is off. No dump hooks. EncoderBuild build_encoder_graph_batched(ggml_context * ctx, const VoxtralWeights & weights, const VoxtralHParams & hp, diff --git a/src/arch/voxtral/model.cpp b/src/arch/voxtral/model.cpp index e1dbc267..b05b8251 100644 --- a/src/arch/voxtral/model.cpp +++ b/src/arch/voxtral/model.cpp @@ -1,15 +1,14 @@ // arch/voxtral/model.cpp - Voxtral (2507) family handler. // // audio-llm: Whisper-large-v3 encoder + 4x frame-group projector + a -// Llama/Ministral causal LM with audio-token injection. Two prompt -// modes share the same encoder/projector/decoder: +// Llama/Ministral causal LM with audio-token injection. Two prompt modes share +// the same encoder/projector/decoder: // // transcription : [BOS][INST][BEGIN_AUDIO][AUDIO]*N[/INST](lang:)?[TRANSCRIBE] // instruct : [BOS][INST][BEGIN_AUDIO][AUDIO]*N BPE(instruction)[/INST] // -// The instruct mode covers translation (task TRANSLATE synthesizes -// "Translate this to {Language}.") and free-text prompting; both are -// mistral-common instruct requests, not transcription requests. +// Instruct mode covers translation (task TRANSLATE synthesizes "Translate this +// to {Language}.") and free-text prompting. #include "voxtral.h" @@ -94,11 +93,10 @@ constexpr const char k_default_variant[] = "voxtral-mini-3b-2507"; constexpr int k_decode_budget_min = 448; // Decode budget (max new text tokens) for an utterance with `n_audio` audio -// embedding tokens. Real speech yields fewer text tokens than audio frames, so -// the audio token count is a safe upper bound; clamp to the context remaining -// under the decoder's trained max so prompt+decode always fits. Greedy decode -// stops at EOS well before this, so a generous ceiling costs nothing but the -// KV allocation it implies. +// embedding tokens. Speech yields fewer text tokens than audio frames, so the +// audio token count is a safe upper bound; clamp to the context remaining under +// the trained max so prompt+decode fits. Greedy decode stops at EOS well before +// this, so a generous ceiling costs only its KV allocation. int pick_decode_budget(int n_audio, int t_prompt, int model_max) { int budget = std::max(k_decode_budget_min, n_audio); const int room = model_max - t_prompt; @@ -106,12 +104,10 @@ int pick_decode_budget(int n_audio, int t_prompt, int model_max) { return budget; } -// Pick a KV-cache context length that fits `needed` (prompt + decode budget) -// tokens. Rounds up to a power of two (so the step graph's `max_n_kv`, computed -// the same way, never exceeds the allocation) and clamps to the decoder's -// trained max_position_embeddings — the real ceiling, since RoPE is only valid -// there. The model's own context (131072 for Voxtral-Mini-3B) is far larger -// than the conservative initial allocation, so long audio grows it on demand. +// Pick a KV-cache context length that fits `needed` (prompt + decode budget). +// Rounds up to a power of two (so the step graph's `max_n_kv`, computed the same +// way, never exceeds the allocation) and clamps to the trained +// max_position_embeddings — the real ceiling, since RoPE is only valid there. int pick_kv_ctx(int needed, int model_max) { int want = 1024; while (want < needed) want *= 2; @@ -119,23 +115,15 @@ int pick_kv_ctx(int needed, int model_max) { return want; } -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Voxtral is a -// hard-context-cap family: the audio tokens, the instruct/transcription -// prompt, and the generated transcript all share the Llama/Ministral decoder's -// context window (dec_max_position_embeddings). The Whisper encoder already -// chunks audio into 30 s windows (375 audio tokens per chunk via the 4x -// projector downsample), so the binding limit is the decoder context, not the -// encoder. Over-length input is rejected up front with -// TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that fills the generation budget -// before end-of-stream is flagged via transcribe_was_truncated(). -// --------------------------------------------------------------------------- +// Input-length contract (see docs/input-limits.md). Hard-context-cap family: +// audio tokens + prompt + transcript share the decoder context window +// (dec_max_position_embeddings); the Whisper encoder chunks audio into 30 s +// windows, so the decoder context is the binding limit. Over-length input is +// rejected with TRANSCRIBE_ERR_INPUT_TOO_LONG; a transcript that fills the +// generation budget before EOS is flagged via transcribe_was_truncated(). -// Generation room the up-front gate reserves, in tokens. Mirrors the floor of -// pick_decode_budget(): an accepted clip is always guaranteed at least this -// many output tokens within the context window. (Longer audio raises the -// actual decode budget above this floor, but the gate only needs to guarantee -// the minimum useful transcript fits.) +// Generation room the up-front gate reserves, in tokens (the pick_decode_budget +// floor): an accepted clip is always guaranteed at least this many output tokens. constexpr int k_gen_reserve = k_decode_budget_min; // Effective decoder context ceiling, in tokens: the model's trained maximum, @@ -149,14 +137,10 @@ int voxtral_context_ceiling(int32_t n_ctx_knob, const VoxtralHParams & hp) { } // Advisory transcribe_capabilities::max_audio_ms: the longest audio whose audio -// tokens, a representative prompt, and the generation reserve still fit the -// context ceiling. The encoder rate is fixed by the 30 s chunk geometry: -// audio_tokens_per_chunk() tokens (375) cover one fe_nb_max_frames-frame chunk -// (3000 mel frames = 30 s), so ms-per-audio-token = chunk_ms / tokens_per_chunk. -// Returns 0 ("unknown / unbounded") if the rate hparams are missing, so a -// misconfigured model is never advertised with a wrong finite number. Note: -// even within this bound a long transcript may truncate at the generation -// budget (transcribe_was_truncated) — max_audio_ms is the input bound. +// tokens + prompt + generation reserve still fit the context ceiling. The +// encoder rate is fixed by the 30 s chunk geometry (ms-per-audio-token = +// chunk_ms / tokens_per_chunk). Returns 0 (unbounded) if the rate hparams are +// missing, so a misconfigured model is never advertised with a wrong number. int64_t voxtral_max_audio_ms(const VoxtralHParams & hp) { const int tokens_per_chunk = hp.audio_tokens_per_chunk(); if (tokens_per_chunk <= 0 || hp.fe_nb_max_frames <= 0 || @@ -321,18 +305,14 @@ transcribe_status load( if (const transcribe_status st = read_voxtral_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) return st; - // Publish the input-length ceiling now that the decoder context window - // and the 30 s-chunk encoder rate are known (apply_family_invariants ran - // before the hparams were read, so it could not set this). See - // docs/input-limits.md. + // Publish the input-length ceiling now that the decoder context window and + // the 30 s-chunk encoder rate are known (apply_family_invariants ran before + // hparams). See docs/input-limits.md. m->caps.max_audio_ms = voxtral_max_audio_ms(m->hparams); - // Basis for the session-level limits query (transcribe_session_get_limits): - // the same constants voxtral_max_audio_ms uses, kept so the effective limit - // can be recomputed at a lowered session n_ctx. The decoder context window - // is the binding cap (audio tokens + prompt + transcript share it), so the - // n_ctx knob and KV-byte estimate agree with caps.max_audio_ms at n_ctx=0. - // See docs/input-limits.md. + // Basis for transcribe_session_get_limits: the same constants + // voxtral_max_audio_ms uses, kept so the effective limit can be recomputed at + // a lowered session n_ctx. { const int tokens_per_chunk = m->hparams.audio_tokens_per_chunk(); if (m->hparams.dec_max_position_embeddings > 0 && tokens_per_chunk > 0 && @@ -484,10 +464,8 @@ transcribe_status init_context( cc->kv_type = params->kv_type; cc->n_ctx = transcribe_session_params_n_ctx(params); - // Encoder is the Whisper-large-v3 encoder (head_dim 64, full - // bidirectional attention, null mask) — flash-supported on every - // backend we ship, matching the whisper arch which runs encoder - // flash unconditionally. Decoder (Llama head_dim 128) also flash. + // Whisper-large-v3 encoder (head_dim 64, full bidirectional, null mask) — + // flash-supported on every backend. Decoder (Llama head_dim 128) also flash. cc->encoder_use_flash = true; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); @@ -705,14 +683,10 @@ transcribe_status run( const int T_prompt = static_cast(prompt_ids.size()); // ----- Input-length gate (see docs/input-limits.md) ----- - // Auto-size the KV cache to this utterance. The initial allocation from - // init_context is a conservative starting point; long audio produces a - // longer audio prompt than that, so grow the cache to fit (capped at the - // context ceiling). The audio-token count is fixed by the input length, so - // reject an over-length clip here, before prefill/decode, instead of - // growing the cache past the trained context and aliasing RoPE. Reserving - // the generation budget means an accepted clip always has room for a real - // transcript. The ceiling honors the caller's n_ctx knob (lowered, never + // Auto-size the KV cache to this utterance (grow to fit, capped at the + // ceiling). The audio-token count is fixed by the input length, so reject an + // over-length clip up front rather than growing past the trained context and + // aliasing RoPE. The ceiling honors the caller's n_ctx knob (lowered, never // raised). const int model_max = voxtral_context_ceiling(cc->n_ctx, cm->hparams); if (T_prompt + k_gen_reserve > model_max) { @@ -905,9 +879,8 @@ transcribe_status run( static_cast(gs)); return TRANSCRIBE_ERR_GGUF; } - // Mid-generation logits dump (validate.py coverage): the step at - // cur_past == T_prompt + 7 produces the reference's scores[8], i.e. - // the logits for the 9th generated token after 8 greedy steps. + // Mid-generation logits dump: the step at cur_past == T_prompt + 7 is + // the reference's scores[8] (logits for the 9th generated token). if (dumps_on && sb.logits != nullptr && cur_past == T_prompt + 7) { transcribe::debug::dump_tensor("dec.logits_raw.gen8", sb.logits, "dec.step"); } @@ -922,10 +895,8 @@ transcribe_status run( } cc->t_decode_us = ggml_time_us() - t_dec_start; - // The decode stopped at EOS (complete) or at the generation budget / - // context width (truncated). Surface the latter via - // transcribe_was_truncated() and a WARN rather than returning a silently - // shortened transcript. See docs/input-limits.md. + // Decode stopped at EOS (complete) or at the generation budget / context + // width (truncated). Surface the latter via transcribe_was_truncated() + WARN. if (next_tok != eos_id) { cc->was_truncated = true; transcribe::log_msg( @@ -959,11 +930,8 @@ transcribe_status run( / static_cast(cm->hparams.fe_sample_rate); cc->segments.push_back(std::move(seg)); - // Output truncation is a hard status: when the single-shot decode stopped - // before EOS (the post-loop check above set cc->was_truncated) the partial - // transcript is already populated and stays readable (like an aborted run); - // surface the truncation to the caller rather than reporting a clean OK. - // See docs/input-limits.md. + // Output truncation is a hard status: the partial transcript stays readable + // (like an aborted run) but the caller is told, not given a clean OK. return cc->was_truncated ? TRANSCRIBE_ERR_OUTPUT_TRUNCATED : TRANSCRIBE_OK; } @@ -1226,10 +1194,8 @@ transcribe_status run_batch( } T_audio_max = std::max(1, T_audio_max); - // Size the batched KV cache to the longest prompt in the batch plus the - // decode budget, clamped to the decoder context ceiling (the trained max, - // lowered by the caller's n_ctx knob — not the single-shot cache's initial - // allocation). kv_init_batched below grows the cache to this on demand. + // Size the batched KV cache to the longest prompt plus the decode budget, + // clamped to the context ceiling. kv_init_batched grows the cache on demand. const int model_max = ctx_ceiling; const int max_new = pick_decode_budget(T_audio_max, max_T_prompt, model_max); int max_n_kv = 1024; @@ -1285,9 +1251,9 @@ transcribe_status run_batch( } std::vector ids(static_cast(max_T_prompt) * n, 0); - // Audio injection by elementwise blend: audio_dense holds each - // utterance's audio embeds scattered into their prompt positions (zero - // elsewhere), keep is 0 there and 1 elsewhere. x = x*keep + audio_dense. + // Audio-injection blend buffers: audio_dense holds each utterance's audio + // embeds at their prompt positions (zero elsewhere); keep is 0 there, 1 + // elsewhere. x = x*keep + audio_dense. std::vector audio_dense(static_cast(dec_h) * max_T_prompt * n, 0.0f); std::vector keep(static_cast(max_T_prompt) * n, 1.0f); std::vector kidx(static_cast(max_T_prompt) * n, 0); diff --git a/src/arch/voxtral/voxtral.h b/src/arch/voxtral/voxtral.h index 1d097939..264ac48c 100644 --- a/src/arch/voxtral/voxtral.h +++ b/src/arch/voxtral/voxtral.h @@ -1,15 +1,11 @@ // arch/voxtral/voxtral.h - Voxtral (2507) model and context types. // -// INTERNAL to src/arch/voxtral/. Declares the concrete classes that -// derive from transcribe_model / transcribe_session for the Voxtral -// audio-LLM family (Whisper-large-v3 encoder + 4x frame-group projector -// + Llama/Ministral causal LM with audio-token injection). -// -// First port targets non-streaming ASR transcription plus a chat -// instruction path (translation / free-text prompting): translation and -// audio understanding both go through the mistral-common instruct -// template, which is structurally distinct from the transcription- -// request template. Both share the same encoder/projector/decoder. +// INTERNAL to src/arch/voxtral/. Concrete transcribe_model / +// transcribe_session classes for the Voxtral audio-LLM family (Whisper-large-v3 +// encoder + 4x frame-group projector + Llama/Ministral causal LM with +// audio-token injection). Two prompt modes — ASR transcription and a +// mistral-common instruct path (translation / free-text) — share the same +// encoder/projector/decoder. #pragma once @@ -42,9 +38,8 @@ namespace transcribe::voxtral { void apply_family_invariants(transcribe_model & model); -// mistral-common transcription/instruct control tokens, resolved through -// the loaded tokenizer at load time so a vocab reorder fails loudly here -// rather than silently corrupting the prompt. +// mistral-common transcription/instruct control tokens, resolved through the +// loaded tokenizer at load time (a vocab reorder fails loudly here). // // Transcription template: // [BOS] [INST] [BEGIN_AUDIO] [AUDIO]*N [/INST] (lang:)? [TRANSCRIBE] @@ -95,7 +90,6 @@ struct VoxtralSession final : public transcribe_session { std::vector mel_buf; std::vector enc_host; // projector output (audio embeds), all chunks - bool encoder_use_flash = false; bool decoder_use_flash = true; diff --git a/src/arch/voxtral/weights.h b/src/arch/voxtral/weights.h index 6a778f8c..5d1776bb 100644 --- a/src/arch/voxtral/weights.h +++ b/src/arch/voxtral/weights.h @@ -1,17 +1,9 @@ // arch/voxtral/weights.h - Voxtral (2507) tensor catalog and hparams. // -// INTERNAL to src/arch/voxtral/. Defines: -// -// - VoxtralHParams: architecture KV that drives tensor shapes. Read -// from stt.voxtral.* and stt.frontend.* at load time, before any -// tensors are allocated. -// -// - VoxtralWeights: named borrowed ggml_tensor* slots. The tensors -// themselves live in the model's ctx_meta / backend buffer. -// -// Architecture pattern: audio-llm (Whisper-large-v3 encoder + 4x frame -// -group projector + Llama/Ministral causal LM with audio-token -// injection). Three-sided weight layout: +// INTERNAL to src/arch/voxtral/. VoxtralHParams: architecture KV (from +// stt.voxtral.* / stt.frontend.*, read before any tensor is allocated). +// VoxtralWeights: named borrowed ggml_tensor* slots (the tensors live in the +// model's ctx_meta / backend buffer). Three-sided weight layout: // // enc.* audio encoder (2x Conv1d stem + 32 bidirectional pre-LN // transformer layers, LayerNorm-with-bias, diff --git a/src/arch/voxtral_realtime/decoder.cpp b/src/arch/voxtral_realtime/decoder.cpp index d68fc7f1..a950d795 100644 --- a/src/arch/voxtral_realtime/decoder.cpp +++ b/src/arch/voxtral_realtime/decoder.cpp @@ -25,7 +25,7 @@ ggml_tensor * named(ggml_tensor * t, const char * name) { } // `ffn_norm_folded` is the per-layer (1+ada(t_cond)) ⊙ norm_ffn_w precomputed per -// session (Stage 6 #2). Passing it as the FFN-norm weight lets the fused +// session. Passing it as the FFN-norm weight lets the fused // rms_norm(·weight) apply the ada scale — no separate per-layer ggml_mul. Null // falls back to the raw FFN-norm weight (ada not yet computed). causal_lm::BlockView to_block_view(const DecBlock & b, ggml_tensor * ffn_norm_folded) { @@ -40,7 +40,7 @@ causal_lm::BlockView to_block_view(const DecBlock & b, ggml_tensor * ffn_norm_fo v.attn_k_norm = nullptr; v.ffn_gate_up_w = b.ffn_gate_up_w; v.ffn_down_w = b.ffn_down_w; - v.ffn_scale = nullptr; // ada scale folded into norm_ffn_w (Stage 6 #2) + v.ffn_scale = nullptr; // ada scale folded into norm_ffn_w return v; } diff --git a/src/arch/voxtral_realtime/decoder.h b/src/arch/voxtral_realtime/decoder.h index ec5e47a4..ba277e46 100644 --- a/src/arch/voxtral_realtime/decoder.h +++ b/src/arch/voxtral_realtime/decoder.h @@ -1,15 +1,13 @@ // arch/voxtral_realtime/decoder.h - Ministral LM prefill/step builders. // -// The per-block math is the shared causal_lm module (GQA, NEOX RoPE, SwiGLU, -// KV cache) with null Q/K-norm and a per-layer FFN-branch scale (ffn_scale) -// carrying the delay-conditioned adaptive-norm `1 + ada(t_cond)`. This file -// owns graph allocation, ADDITIVE audio fusion (inputs_embeds += audio at -// every position), the TIED lm_head, and dump naming. +// The per-block math is the shared causal_lm module (GQA, NEOX RoPE, SwiGLU, KV +// cache) with null Q/K-norm and a per-layer FFN-branch scale (ffn_scale) +// carrying the delay-conditioned adaptive-norm `1 + ada(t_cond)`. This file owns +// graph allocation, ADDITIVE audio fusion, the TIED lm_head, and dump naming. // // Audio fusion: x[:, i] = embed_tokens(id[i]) + audio[:, i] for all i. The -// caller supplies exactly T audio embeddings ([dec_hidden, T]); the prefill -// over the prompt uses the first T_prompt audio embeds, each decode step adds -// the audio embed for its position. +// caller supplies exactly T audio embeddings ([dec_hidden, T]); each decode step +// adds the audio embed for its position. #pragma once @@ -79,15 +77,13 @@ StepBuild build_step_graph(ggml_context * ctx, bool use_flash); // --------------------------------------------------------------------------- -// Multi-position verify (spec-decode prototype) +// Multi-position verify (speculative decode) // --------------------------------------------------------------------------- // // Runs the decoder on T_verify positions in one pass. Used by an n-gram // speculative-decode driver: feed [next_tok, draft[0..K-1]] (T_verify = K+1), // get back T_verify predicted tokens, accept the longest prefix where draft[i] -// matches the predicted token at column i, rewind by emitting up to that -// point. KV writes are committed at positions kv_idx[T_verify]; if some -// draft tokens are rejected, the wrongly-written KV rows are simply +// matches the predicted token at column i. Rejected drafts' KV rows are simply // overwritten on the next pass (positions are addressed by slot). struct VerifyBuild { ggml_tensor * input_ids_in = nullptr; // [T_verify] i32 @@ -118,11 +114,10 @@ VerifyBuild build_verify_graph(ggml_context * ctx, // Mirrors the single-utterance builders with the batch on ne[2], via the shared // causal_lm::block_{prefill,step}_batched (flash-only, n_batch=B kv cache). The // realtime specifics ride along: TIED lm_head, the per-layer ada ffn_scale view, -// and ADDITIVE audio fusion. The prompt is uniform length across the batch (BOS -// + STREAMING_PAD*(32+delay) with `num_delay` shared), so the prefill is -// rectangular [hidden, T_prompt, B] — no ragged padding; the last real position -// is T_prompt-1 for every row. Per-step audio injection (a fresh audio embed per -// position) is why the decode loop cannot reuse causal_lm::run_batched_step_loop. +// and ADDITIVE audio fusion. The prompt is uniform length across the batch, so +// the prefill is rectangular [hidden, T_prompt, B] (last real position T_prompt-1 +// for every row). Per-step audio injection is why the decode loop cannot reuse +// causal_lm::run_batched_step_loop. struct PrefillBuildBatched { ggml_tensor * input_ids_in = nullptr; // [T_prompt, B] i32 diff --git a/src/arch/voxtral_realtime/encoder.cpp b/src/arch/voxtral_realtime/encoder.cpp index 13da067c..a24b9e52 100644 --- a/src/arch/voxtral_realtime/encoder.cpp +++ b/src/arch/voxtral_realtime/encoder.cpp @@ -106,9 +106,8 @@ ggml_tensor * mha(ggml_context * ctx, ggml_tensor * x, // Batched full-MHA self-attention: x:[d_model, T, B]. `positions [T]` and the // sw-causal `mask [T, T]` are SHARED across the batch (the encoder is causal, so -// right-pad rows are numerically isolated and need no per-row mask). Flash-only. -// Mirrors `mha` with the batch riding ne[2] of x (and ne[3] of the per-head -// tensors), exactly like causal_lm::block_prefill_batched on the decoder side. +// right-pad rows are numerically isolated and need no per-row mask). Mirrors +// `mha` with the batch riding ne[2] of x (ne[3] of the per-head tensors). ggml_tensor * mha_batched(ggml_context * ctx, ggml_tensor * x, const EncBlock & b, ggml_tensor * positions, ggml_tensor * mask, int n_heads, int head_dim, int /*d_model*/, float rope_theta, diff --git a/src/arch/voxtral_realtime/encoder.h b/src/arch/voxtral_realtime/encoder.h index 4b92c515..5765aa66 100644 --- a/src/arch/voxtral_realtime/encoder.h +++ b/src/arch/voxtral_realtime/encoder.h @@ -58,16 +58,14 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, // Batched offline encoder (transcribe_run_batch fast path) // --------------------------------------------------------------------------- // -// Mirrors the reference, which runs the audio tower batched over [B, T, C] with -// NO audio attention_mask: every clip's mel is RIGHT-padded to the batch-max -// `n_mel_frames` and stacked on ne[2]=B. Because the conv stem is causal -// (left-pad) and attention is causal + sliding-window, a real frame at position -// t only attends to <= t, so the right-padding of shorter clips never -// contaminates their real [0, T_enc_b) outputs. A SINGLE `positions [T_enc]` -// and sw-causal `mask [T_enc, T_enc]` therefore serve every row; the caller -// slices each row's real n_audio_b out of out[:, :, b]. `use_flash` selects the -// flash vs manual-softmax attention (honor the session's encoder_use_flash, so -// the default CPU source-of-truth non-flash path batches too). +// Runs the audio tower batched over [B, T, C] with NO audio attention_mask: +// every clip's mel is RIGHT-padded to the batch-max `n_mel_frames` and stacked on +// ne[2]=B. The conv stem is causal (left-pad) and attention is causal + +// sliding-window, so a real frame at t only attends to <= t — the right-padding +// of shorter clips never contaminates their real [0, T_enc_b) outputs. A SINGLE +// `positions [T_enc]` and sw-causal `mask [T_enc, T_enc]` serve every row; the +// caller slices each row's real n_audio_b out of out[:, :, b]. `use_flash` +// honors the session's encoder_use_flash (so the non-flash CPU path batches too). struct EncoderBuildBatched { ggml_tensor * mel_in = nullptr; // [n_mels, n_mel_frames, B] input ggml_tensor * positions_in = nullptr; // [T_enc] i32 RoPE positions (shared) @@ -91,22 +89,16 @@ EncoderBuildBatched build_encoder_graph_batched(ggml_context * ctx, // Incremental streaming encoder (matches the reference StaticCache mechanism) // --------------------------------------------------------------------------- // -// The reference runs the encoder transformer INCREMENTALLY: each decode step -// feeds `downsample_factor` new embedder frames through the 32 transformer -// layers against an encoder StaticCache (sliding_window=750), producing one -// audio token. We split that into two graphs: +// The encoder transformer runs INCREMENTALLY against an encoder StaticCache +// (sliding_window=750). Split into two graphs: // // build_embedder_graph — conv stem only: mel -> embedder frames // [d_model, T_enc]. Causal convs (left-pad), so the -// frames are append-only/stable (a whole-buffer -// recompute equals the reference per-chunk conv + -// padding cache, frame-for-frame). -// build_encoder_chunk_graph — the expensive part, INCREMENTAL: `n_new` -// embedder frames as queries attend to the encoder -// KV cache (prior frames' K/V) under a sliding- -// window-causal mask; their K/V are appended at -// cache offset `n_kv_before`. Final RMSNorm + -// projector group-of-4 -> audio embeds. +// frames are append-only/stable. +// build_encoder_chunk_graph — INCREMENTAL: `n_new` embedder frames as queries +// attend to the encoder KV cache under a sliding- +// window-causal mask; their K/V are appended at the +// cache offset. Final RMSNorm + projector group-of-4. struct EmbedderBuild { ggml_tensor * mel_in = nullptr; // [n_mels, n_mel_frames] input @@ -121,14 +113,12 @@ EmbedderBuild build_embedder_graph(ggml_context * ctx, const HParams & hp, int n_mel_frames); -// Conv stem with the streaming PADDING CACHE (reference VoxtralRealtimeConv1dCacheLayer). -// Feeds only `n_new_mel` NEW mel frames; the conv left-context comes from the host- -// carried caches instead of zero left-pad: conv1 (k3 s1, left_pad 2) ← last 2 mel -// frames; conv2 (k3 s2, left_pad 1) ← last 1 conv1-output frame. Both zero on the -// first chunk, which makes the cached path bit-identical to the whole-buffer left-pad -// conv. Emits M_emb = n_new_mel/2 new embedder frames (n_new_mel must be even) plus -// `cache2_out` = the chunk's last conv1-output frame for the next call. The conv1 -// cache (last 2 mel frames) is trivially the input tail, carried host-side. +// Conv stem with the streaming PADDING CACHE. Feeds only `n_new_mel` NEW mel +// frames; the conv left-context comes from host-carried caches instead of zero +// left-pad: conv1 (k3 s1, left_pad 2) ← last 2 mel frames; conv2 (k3 s2, left_pad +// 1) ← last 1 conv1-output frame. Both zero on the first chunk == whole-buffer +// left-pad. Emits M_emb = n_new_mel/2 new embedder frames (n_new_mel must be even) +// plus `cache2_out` = the chunk's last conv1-output frame for the next call. struct EmbedderChunkBuild { ggml_tensor * mel_in = nullptr; // [n_mels, n_new_mel] new mel frames ggml_tensor * cache1_in = nullptr; // [n_mels, 2] prev 2 mel frames (conv1 left ctx) diff --git a/src/arch/voxtral_realtime/model.cpp b/src/arch/voxtral_realtime/model.cpp index 583d07b3..f05414a7 100644 --- a/src/arch/voxtral_realtime/model.cpp +++ b/src/arch/voxtral_realtime/model.cpp @@ -1,14 +1,10 @@ // arch/voxtral_realtime/model.cpp - Voxtral Realtime (2602) family handler. // -// Streaming audio-LLM. Stage 4 implements the OFFLINE whole-clip forward (the -// per-tensor `decode` oracle contract): causal RoPE sliding-window encoder -> -// projector -> Ministral LM with ADDITIVE audio fusion and delay-conditioned -// adaptive-norm. The prompt is [BOS] + [STREAMING_PAD]*(n_left_pad+num_delay); -// the projector audio embeddings are ADDED onto every sequence position. Greedy -// decode is clamped to ceil(mel_frames / audio_length_per_tok) positions. -// -// The incremental streaming scheduler (conv cache + dual KV + chunk loop + -// stream hooks) lands in a later step; this file ships the offline baseline. +// Streaming audio-LLM: causal RoPE sliding-window encoder -> projector -> +// Ministral LM with ADDITIVE audio fusion and delay-conditioned adaptive-norm. +// The prompt is [BOS] + [STREAMING_PAD]*(n_left_pad+num_delay); the projector +// audio embeddings are ADDED onto every sequence position. Greedy decode is +// clamped to ceil(mel_frames / audio_length_per_tok) positions. #include "voxtral_realtime.h" #include "transcribe/voxtral_realtime.h" // public stream-ext surface @@ -123,59 +119,24 @@ bool read_ref_mel(const std::string & dir, int n_mels, return true; } -// --------------------------------------------------------------------------- -// Input-length contract (see docs/input-limits.md). Voxtral Realtime is the -// honest edge case among the audio-LLM families: BOTH ends are effectively -// unbounded. -// -// - Encoder: causal + sliding-window (750 frames). The streaming scheduler -// advances it frame-by-frame against a constant-memory ring, so input -// length is bounded only by wall-clock, not by an architectural wall. -// - Decoder: a Ministral LM whose KV is itself a sliding-window ring -// (dec_sliding_window = 8192) keeping the last `swin` tokens. RoPE stays -// valid up to dec_max_position (131072), so the only hard cap is that the -// decoder position cannot exceed dec_max_position. -// -// dec_max_position = 131072 positions at the 12.5 Hz audio-token grid -// (audio_length_per_tok = 8 mel frames/token, hop 1280 samples/token @ 16 kHz -// => 80 ms/token) is 131072 * 80 ms ~= 2.9 hours of audio before the absolute -// position cap is reached. That is far past the ~1 hour threshold in -// docs/input-limits.md AND the family is genuinely stream-unbounded (the -// decoder KV slides, it does not clamp short of dec_max_position), so the -// honest published value is 0 (unbounded). We do NOT advertise a finite -// max_audio_ms: the real limit a streaming caller hits first is memory/latency, -// not tokens, and a finite advisory number would imply a reject at an everyday -// threshold we don't enforce. There IS a genuine hard wall — the absolute -// dec_max_position cap (~2.9 h), past which RoPE positions alias — and the -// offline paths DO reject beyond it: one-shot run() and run_batch return -// INPUT_TOO_LONG, while streaming surfaces it after the fact via -// transcribe_was_truncated() + a WARN (a stream is never forced to FAILED). -// But that wall is hours out, far past any practical clip, so 0 (unbounded) -// stays the honest capability and n_ctx does not lower it. -// --------------------------------------------------------------------------- - -// Absolute decoder position cap, in tokens: the model's maximum RoPE position -// (dec_max_position). Past this the decoder's RoPE positions alias and the -// decode is invalid, so every decode path (one-shot / batch / streaming) -// hard-stops here. This is the model's true wall and is deliberately NOT -// lowered by the session n_ctx knob: voxtral_realtime is a chunked/unbounded -// (bucket-1) family whose decoder KV slides, so n_ctx neither bounds its (ring) -// memory nor changes this cap. Like whisper and parakeet, the family ignores -// n_ctx entirely — transcribe_session_get_limits() reports it unbounded and a -// lowered n_ctx is a documented no-op. See the contract block above and -// docs/input-limits.md. +// Input-length contract (see docs/input-limits.md). Chunked/unbounded family: +// the encoder is causal + sliding-window (750 frames) advanced on a +// constant-memory ring, and the decoder KV is itself a sliding-window ring +// (dec_sliding_window = 8192). The only hard wall is the dec_max_position +// (131072) RoPE cap; at the 80 ms/token grid (audio_length_per_tok=8 mel +// frames, hop 1280 @ 16 kHz) that is ~2.9 h. n_ctx does not lower it. + +// Absolute decoder position cap, in tokens (dec_max_position). Past this the +// decoder's RoPE positions alias, so every decode path hard-stops here. Not +// lowered by the session n_ctx knob (chunked/unbounded family; the decoder KV +// slides). See docs/input-limits.md. int voxtral_realtime_abs_position_cap(const HParams & hp) { return hp.dec_max_position; } // Advisory transcribe_capabilities::max_audio_ms. Returns 0 (unbounded): the -// decoder KV is a sliding-window ring (input length is bounded only by the -// dec_max_position absolute-position cap, ~2.9 h at the 80 ms/token grid) and -// the encoder is causal/streaming, so this family is in the chunked/unbounded -// bucket. The honest number for "no practical limit" is 0; a finite advertised -// value would imply an everyday caller-facing limit. The only reject is the -// model's true absolute position wall, which is hours out and handled as a -// safety cap. See the contract block above and docs/input-limits.md. +// decoder KV slides and the encoder is causal/streaming, so the only reject is +// the model's absolute position wall (~2.9 h), handled as a safety cap. int64_t voxtral_realtime_max_audio_ms(const HParams & hp) { (void) hp; return 0; @@ -204,10 +165,8 @@ transcribe_status load(Loader & loader, if (auto st = read_hparams(loader.gguf(), m->hparams); st != TRANSCRIBE_OK) return st; // Publish the input-length ceiling now that the decoder context window is - // known (apply_family_invariants ran before the hparams were read, so it - // could not set this). Voxtral Realtime is chunked/unbounded: 0 = no - // practical limit (sliding-window KV + streaming encoder). See - // docs/input-limits.md and the contract block in the anon namespace. + // known (apply_family_invariants ran before hparams). Chunked/unbounded: + // 0 = no practical limit. See docs/input-limits.md. m->caps.max_audio_ms = voxtral_realtime_max_audio_ms(m->hparams); m->hparams.vocab_size = m->tok.n_tokens(); @@ -320,19 +279,14 @@ transcribe_status init_context(transcribe_model * model, cc->model = model; cc->n_threads = params->n_threads; cc->kv_type = params->kv_type; - // Captured for base-class uniformity only: voxtral_realtime is a - // chunked/unbounded family and deliberately ignores n_ctx (its decoder KV - // ring is constant-memory; the decode wall is the absolute dec_max_position, - // not a lowerable ceiling). See docs/input-limits.md. + // Captured for base-class uniformity only: this family ignores n_ctx (the + // 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 (Stage 6 - // #3 — measured ~14-37% faster encode on Metal, ~17-47% on CPU, byte-identical - // transcript). Flash is now the numerical source-of-truth for the encoder; - // tests/tolerances/voxtral_realtime.json is gated against it. Override with - // TRANSCRIBE_NO_FLASH=1. + // 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); @@ -443,10 +397,10 @@ transcribe_status compute_ada_scales(Session * cc, Model * cm, int num_delay) { ggml_backend_tensor_get(acc, ada.data(), 0, ada.size() * sizeof(float)); ggml_free(ctx); for (auto & v : ada) v += 1.0f; // s_l = 1 + ada - // Stage 6 #2: fold the per-layer FFN-norm weight into the ada scale and store - // (s_l ⊙ norm_ffn_w) here. The decode block passes this as the FFN-norm weight - // (ffn_scale = null), so the fused rms_norm(·weight) kernel applies the ada - // scale with NO separate per-layer ggml_mul. Exact: s⊙(w⊙x̂) == (s⊙w)⊙x̂. + // Fold the per-layer FFN-norm weight into the ada scale and store + // (s_l ⊙ norm_ffn_w): the decode block passes this as the FFN-norm weight so + // the fused rms_norm(·weight) kernel applies the ada scale with no separate + // ggml_mul. Exact: s⊙(w⊙x̂) == (s⊙w)⊙x̂. { std::vector wbuf(static_cast(hidden)); for (int il = 0; il < n_layer; ++il) { @@ -493,7 +447,6 @@ transcribe_status forward_buffer(Session * cc, Model * cm, // [n_left_pad zeros] + [raw rounded up to raw_per_tok] + [n_right zeros] // raw_per_tok = audio_length_per_tok * hop = 1280 (= sr / frame_rate); // n_right = num_delay + 1 (BOS) + OFFLINE_STREAMING_BUFFER_TOKENS(10). - // This reproduces the reference input_features exactly (max|d|=0). const int raw_per_tok = cm->hparams.audio_length_per_tok * cm->hparams.fe_hop_length; const int n_left_tok = cm->specials.n_left_pad; const int n_right_tok = num_delay + 1 + 10; @@ -624,12 +577,9 @@ transcribe_status forward_buffer(Session * cc, Model * cm, return TRANSCRIBE_ERR_INVALID_ARG; } - // Absolute-position consistency check (parity with the batch path): the - // decoder's RoPE is valid only up to dec_max_position, and a clip needs one - // KV slot per audio token, so a horizon past the cap cannot be decoded. - // Reject up front with INPUT_TOO_LONG. This family ignores the session - // n_ctx knob, so the cap is always the model's true wall (~2.9 h); a clip - // that long OOMs the grow-to-fit KV first, so this is a defensive guard. + // RoPE is valid only up to dec_max_position and a clip needs one KV slot per + // audio token, so a horizon past the cap cannot be decoded — reject up front + // with INPUT_TOO_LONG (defensive: ~2.9 h of audio OOMs the KV first). const int abs_cap = voxtral_realtime_abs_position_cap(cm->hparams); if (n_audio + 1 > abs_cap) { transcribe::log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -716,22 +666,11 @@ transcribe_status forward_buffer(Session * cc, Model * cm, generated_ids.push_back(argmax(logits)); } - // Steps: 1-gram-lookup speculative decode (when k_drafts > 0) or plain - // autoregression (k_drafts == 0). - // - // Spec path: each iter feeds [next_tok, draft[0..K-1]] at positions - // [cur_pos..cur_pos+K] through a multi-position verify graph. predicted[0] - // is always the model's true next-token decision; predicted[c] (c>0) is - // only valid if draft[0..c-1] all matched the model's argmax at columns - // 0..c-1. n_accept = longest matched prefix; we commit predicted[0..n_accept]. - // Rejected drafts' KV writes get overwritten on the next iter (subsequent - // verify pass at cur_pos + n_accept + 1 writes through those slots). - // - // For both paths the attention read window is shrunk from kv_cache.n_ctx - // (allocation stride) to n_audio (the actual clip horizon). The dropped - // slots [n_audio, n_ctx) were always -inf-masked, so the shrink is bit- - // identical and saves ~11× KV bandwidth on short clips where n_ctx floors - // at 2048. Streaming path keeps the full window because positions wrap. + // Steps: 1-gram-lookup speculative decode (k_drafts > 0; mechanism in + // build_verify_graph) or plain autoregression (k_drafts == 0). For both, the + // attention read window is shrunk from kv_cache.n_ctx to n_audio (the actual + // clip horizon): the dropped slots [n_audio, n_ctx) were always -inf-masked, + // so the shrink is bit-identical and saves KV bandwidth on short clips. { const int max_n_kv = n_audio; @@ -1050,18 +989,12 @@ transcribe_status run(transcribe_session * session, const float * pcm, int n_sam // 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 K=1: the robust setting across backends. A metal Q4_K_M sweep - // (decode_ms vs K, jfk + dots) showed K=1 never regresses — best on - // short utterances (~1.12x) and break-even on long ones — while K>=2 - // costs long-form decode because 1-gram draft acceptance is too low to - // amortize the T=K+1 verify graph (K=2 dots ~ -16%, K>=4 far worse). - // The win concentrates on bandwidth-bound hardware, where one extra - // drafted position rides along on weights already loaded for the - // verify step. See docs/models/voxtral-realtime.md. + // misconfigured caller doesn't ask for an unbounded verify graph. Default + // K=1 is the robust setting across backends (K>=2 costs long-form decode; + // 1-gram acceptance is too low to amortize the T=K+1 verify graph). See + // docs/models/voxtral-realtime.md. constexpr int VOXTRAL_REALTIME_SPEC_K_MAX = 8; - int k_drafts = 1; // family default + int k_drafts = 1; if (params != nullptr && params->struct_size >= offsetof(transcribe_run_params, spec_k_drafts) + sizeof(params->spec_k_drafts)) { const int requested = params->spec_k_drafts; @@ -1088,16 +1021,12 @@ transcribe_status run(transcribe_session * session, const float * pcm, int n_sam } // ---- Streaming ------------------------------------------------------------- -// Voxtral Realtime's encoder is causal + sliding-window and its mel uses a -// FIXED global normalization, so the projector audio embeddings are append-only -// and stable: a forward over the accumulated audio yields embeddings identical -// to the offline whole-clip forward. stream_finalize therefore runs the exact -// offline forward on the full buffer => byte-identical to transcribe_run and the -// `stream` oracle. stream_feed emits throttled tentative hypotheses by running -// the same forward on the audio-so-far. (Constant-per-chunk-cost incremental -// encode — conv padding cache + encoder StaticCache + decoder sliding-KV — is a -// documented Stage-5+ perf optimization with identical numerics; see -// reports/porting/voxtral_realtime/streaming-plan.md.) +// The encoder is causal + sliding-window and the mel uses a FIXED global +// normalization, so the projector audio embeddings are append-only and stable: +// a forward over the accumulated audio yields embeddings identical to the +// offline whole-clip forward. The incremental scheduler below (conv padding +// cache + encoder StaticCache ring + decoder sliding-KV) is therefore numerically +// identical to the offline path; stream_finalize matches transcribe_run. constexpr int k_default_min_decode_interval_ms = 1000; @@ -1117,11 +1046,9 @@ transcribe_status resolve_stream_ext(const transcribe_stream_params * sp, if (fam != nullptr) { const auto * vx = reinterpret_cast(fam); - // num_delay_tokens: -1 = model default (6 = 480 ms). Otherwise it must be - // a publisher-supported transcription delay. mistral-common requires a - // positive multiple of 80 ms (audio.py:124-127); the model card narrows - // the validated set to a multiple of 80 ms in [80, 1200] (tokens 1..15) - // OR the standalone 2400 ms (token 30). One token = 80 ms (12.5 Hz grid). + // num_delay_tokens: -1 = model default (6 = 480 ms). Otherwise a + // publisher-validated delay: tokens 1..15 (80..1200 ms) or 30 (2400 ms). + // One token = 80 ms (12.5 Hz grid). const int32_t nt = vx->num_delay_tokens; if (nt != -1 && !((nt >= 1 && nt <= 15) || nt == 30)) return TRANSCRIBE_ERR_INVALID_ARG; @@ -1197,13 +1124,11 @@ bool stream_run_graph(Session * cc, Model * cm, ggml_cgraph * gf, return ok; } -// Encoder KV ring geometry. The reference keeps the last `sliding_window`(750) -// frames (StaticSlidingWindowLayer). We hold a contiguous cache of k_enc_ring_ctx -// slots, append new frames at the write head, and periodically COMPACT (copy the -// last 750 frames to the front) so the ring never overflows — constant memory, -// unbounded stream length, numerically identical to the whole-clip encoder. -// k_enc_max_batch bounds frames per chunk so a single huge feed can't overflow -// (must be <= ring - sliding_window and a multiple of proj_downsample). +// 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; // (decoder sliding-window ring size is read from the GGUF: hp.dec_sliding_window) @@ -1260,10 +1185,9 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * // needs: absolute samples [ws_sample, total_abs). ws_sample backs off // GUARD frames before the first uncommitted mel frame; center=True // reflects at the window start, so the first ceil(pad/hop)=2 window - // frames are discarded (GUARD=4). The window is the SAME samples the - // whole-buffer mel saw -> bit-identical. PCM older than the window is - // dropped (below), so memory is bounded. `mel_off` = absolute frame of - // window-frame 0; `mel_n` = window frame count (stride into mel_buf). ---- + // frames are discarded (GUARD=4). Same samples the whole-buffer mel + // saw -> bit-identical. `mel_off` = absolute frame of window-frame 0; + // `mel_n` = window frame count (stride into mel_buf). ---- constexpr int GUARD = 4; const size_t left = static_cast(n_left_tok) * raw_per_tok; const int mel_off = std::max(0, cc->stream_n_mel_committed - GUARD); @@ -1321,10 +1245,9 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * if (is_final) cc->stream_n_audio_clamp = T_emb_ready / down; // ---- 3. Conv stem with the streaming PADDING CACHE: feed only the NEW mel - // frames (token-aligned, a multiple of 2*down=8) through the conv stem - // against the carried conv caches -> M_emb = M/2 new embedder frames. - // Bit-identical to the whole-buffer left-pad conv, but O(new) per feed - // (reference VoxtralRealtimeConv1dCacheLayer). ---- + // frames (token-aligned, a multiple of 2*down=8) against the carried + // conv caches -> M_emb = M/2 new embedder frames. Bit-identical to the + // whole-buffer left-pad conv, but O(new) per feed. ---- const int ed = hp.enc_d_model; const int M = ((n_mel_ready - cc->stream_n_mel_committed) / (2 * down)) * (2 * down); if (M > 0) { @@ -1535,23 +1458,17 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * const int swin = hp.dec_sliding_window; // 8192 const int kv_ring = max_n_kv; // ring size == n_ctx (== swin here) - // Hard absolute-position cap = dec_max_position (131072). This family - // ignores the session n_ctx knob (bucket-1 / unbounded), so the cap is - // always the model's true wall — a streaming caller hits memory/latency - // long before it (~2.9 h of continuous audio). + // 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); const ggml_fp16_t mz = ggml_fp32_to_fp16(0.0f), mn = ggml_fp32_to_fp16(-INFINITY); std::vector step_mask(max_n_kv, mn); int cap = (is_final && cc->stream_n_audio_clamp >= 0) ? std::min(cc->stream_n_tok_ready, cc->stream_n_audio_clamp) : cc->stream_n_tok_ready; - // The decoder KV is a sliding-window ring (keeps the last `swin` tokens, - // matching the reference DynamicSlidingWindowLayer), so length is bounded - // only by the model's max position — clips past the window slide, not clamp. // `hit_pos_cap`: this chunk's decode was limited by the hard position cap - // (NOT by stream_n_tok_ready / the finalize clamp), i.e. the stream is at - // the absolute context wall. Distinguish it from a normal per-chunk stop - // so the truncation WARN below fires only at the genuine hard cap, once. + // (not by stream_n_tok_ready / the finalize clamp), so the truncation WARN + // below fires only at the genuine hard cap. const bool hit_pos_cap = (max_pos < cap); cap = std::min(cap, max_pos); while (!cc->stream_eos && cc->stream_dec_pos < cap) { @@ -1590,15 +1507,9 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * cc->stream_generated.push_back(tok); } - // Non-silent truncation at the hard context cap. The decode loop above - // clamps stream_dec_pos to dec_max_position (max_pos); when the stream - // reaches that wall WITHOUT an EOS, the transcript is incomplete. Flag it - // via transcribe_was_truncated() + a single WARN rather than letting the - // clamp drop tokens silently. Guarded by !was_truncated so it fires at - // most once even though stream_process is re-entered per feed/finalize. - // The clamp/stopping behavior itself is unchanged. (Normal per-chunk - // stops — dec_pos reaching stream_n_tok_ready or the finalize clamp — - // never set hit_pos_cap, so they do not warn.) + // Truncation at the hard context cap without an EOS: flag via + // transcribe_was_truncated() + a single WARN (guarded by !was_truncated + // so it fires at most once across the per-feed/finalize re-entries). if (hit_pos_cap && !cc->stream_eos && cc->stream_dec_pos >= cap && !cc->was_truncated) { cc->was_truncated = true; @@ -1613,9 +1524,8 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * cc->stream_t_dec_us += ggml_time_us() - t_dec0; // Release consumed audio embeds [audio_base, dec_pos) — the decoder never - // re-reads past positions, so memory stays bounded (matches the reference, - // which consumes each embed at its step). Skipped while dumping so the - // finalize proj.out dump keeps the full [0, n_tok_ready) history. + // re-reads past positions, so memory stays bounded. Skipped while dumping so + // the finalize proj.out dump keeps the full [0, n_tok_ready) history. if (!dumps_on && cc->stream_dec_pos > cc->stream_audio_base) { const size_t drop_n = static_cast(cc->stream_dec_pos - cc->stream_audio_base) * dec_h; cc->stream_audio_embeds.erase(cc->stream_audio_embeds.begin(), @@ -1644,8 +1554,7 @@ transcribe_status stream_process(Session * cc, Model * cm, bool is_final, bool * // ---- 6. Numerical-validation dumps (finalize only). ---- if (dumps_on && is_final) { - // stream proj.out = the incremental-encoder audio embeds [dec_h, n_tok]; - // diff vs offline proj.out proves the StaticCache encoder is bit-exact. + // stream proj.out = the incremental-encoder audio embeds [dec_h, n_tok]. if (cc->stream_n_tok_ready > 0) { const long long shp[2] = { dec_h, cc->stream_n_tok_ready }; transcribe::debug::dump_host_f32("proj.out", cc->stream_audio_embeds.data(), @@ -1687,7 +1596,6 @@ transcribe_status stream_begin(transcribe_session * session, cc->stream_n_enc_committed = 0; cc->stream_enc_slot = 0; cc->stream_enc_abs_base = 0; - // Conv-stem padding cache: zeros == whole-buffer left-pad on the first chunk. 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; @@ -1719,15 +1627,12 @@ transcribe_status stream_begin(transcribe_session * session, if (cc->enc_kv.buffer != nullptr) ggml_backend_buffer_clear(cc->enc_kv.buffer, 0); // Decoder KV: a sliding-window RING sized to the model's own `sliding_window` - // (read from the GGUF, NOT hardcoded — 8192 for this variant). 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 — matching the reference - // DynamicSlidingWindowLayer (keep last swin-1 + new). `sliding_window` is a - // trained-in architectural constant, not an inference knob: shrinking it feeds - // the model an attention span it never saw; growing it does nothing (the model - // only attends `swin` back). Sized once per session; never shrinks a larger - // ctx an offline run may have left. - const int dec_ring = cm->hparams.dec_sliding_window; // == reference text_config.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; 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(); @@ -1839,14 +1744,12 @@ bool accepts_ext_kind(const transcribe_model * model, transcribe_ext_slot slot, // Offline batched decode (transcribe_run_batch) // --------------------------------------------------------------------------- // -// Mirrors the reference's batched audio tower: every clip's mel is right-padded -// to the batch max and stacked on ne[2]; the causal encoder isolates each row's -// real frames (no per-row mask). Batched additive fusion + a lockstep batched -// decoder follow. The realtime specifics vs voxtral's run_batch: a single -// whole-clip encoder (not 30 s chunks), TIED head, the ada ffn_scale, ADDITIVE -// audio injected at EVERY prefill position AND at every decode step, and a -// per-row n_audio clamp. Falls back to serial for n==1, dump mode, or no decoder -// flash (the batched causal_lm blocks are flash-only). +// Batched audio tower: every clip's mel is right-padded to the batch max and +// stacked on ne[2]; the causal encoder isolates each row's real frames (no +// per-row mask). vs voxtral's run_batch: a single whole-clip encoder (not 30 s +// chunks), TIED head, the ada ffn_scale, ADDITIVE audio at every prefill position +// AND every decode step, and a per-row n_audio clamp. Falls back to serial for +// n==1, dump mode, or no decoder flash (the batched blocks are flash-only). transcribe_status run_batch_serial(Session * cc, const float * const * pcm, const int * n_samples, int n, @@ -2085,15 +1988,10 @@ transcribe_status run_batch(transcribe_session * session, const float * const * for (int i = 0; i < sp.n_left_pad + num_delay; ++i) prompt_ids.push_back(sp.streaming_pad); const int T_prompt = static_cast(prompt_ids.size()); - // model_max = dec_max_position (131072), the model's absolute RoPE wall. - // This family ignores the session n_ctx knob (bucket-1 / unbounded), so the - // cap is always the model's true max — n_ctx never narrows it. The realtime - // decoder needs one KV slot per audio token, so a clip whose audio horizon - // exceeds model_max CANNOT be partially decoded by clamping the cache below - // it (that corrupts the decode). Reject it up front with INPUT_TOO_LONG — - // the same contract as the offline run() path. (Reaching this wall needs - // ~2.9 h of audio; a clip that long OOMs the linear batch KV first.) See - // docs/input-limits.md. + // model_max = dec_max_position, the absolute RoPE wall (n_ctx never narrows + // it). One KV slot per audio token, so a clip whose horizon exceeds model_max + // can't be partially decoded by clamping the cache — reject up front with + // INPUT_TOO_LONG, same contract as run(). const int model_max = voxtral_realtime_abs_position_cap(cm->hparams); int max_audio = 0; for (int b = 0; b < n; ++b) { @@ -2119,9 +2017,9 @@ transcribe_status run_batch(transcribe_session * session, const float * const * // pow2 width is safely clamped to the ceiling without cutting any decode. int max_n_kv = 1024; while (max_n_kv < max_audio + 1) max_n_kv *= 2; if (model_max > 0 && max_n_kv > model_max) max_n_kv = model_max; - // Step attention reads only the valid window (Stage 6 #4): the batch's actual - // max audio length, not the power-of-2 allocation. min() preserves the full - // window if the alloc was capped. Bit-identical (extra slots are masked). + // Step attention reads only the valid window: the batch's actual max audio + // length, not the power-of-2 allocation. min() preserves the full window if + // the alloc was capped. Bit-identical (extra slots are masked). const int read_n_kv = std::min(max_audio, max_n_kv); const ggml_type kv_type = (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) ? GGML_TYPE_F32 : GGML_TYPE_F16; if (cc->kv_cache_batch.self_k == nullptr || cc->kv_batch_cap != n || @@ -2301,6 +2199,6 @@ extern "C" void transcribe_voxtral_realtime_stream_ext_init( std::memset(p, 0, sizeof(*p)); p->ext.size = sizeof(*p); p->ext.kind = TRANSCRIBE_EXT_KIND_VOXTRAL_REALTIME_STREAM; - p->num_delay_tokens = -1; // model default - p->min_decode_interval_ms = -1; // family default + p->num_delay_tokens = -1; + p->min_decode_interval_ms = -1; } diff --git a/src/arch/voxtral_realtime/voxtral_realtime.h b/src/arch/voxtral_realtime/voxtral_realtime.h index 47e30bc9..5bff35b8 100644 --- a/src/arch/voxtral_realtime/voxtral_realtime.h +++ b/src/arch/voxtral_realtime/voxtral_realtime.h @@ -3,8 +3,6 @@ // INTERNAL to src/arch/voxtral_realtime/. Streaming audio-LLM: causal RoPE // sliding-window audio encoder + 4x frame-group projector + Ministral causal // LM with ADDITIVE audio fusion and per-layer delay-conditioned adaptive-norm. -// Stage 4 builds the offline whole-clip forward first (the per-tensor -// contract), then the incremental streaming scheduler. #pragma once @@ -90,18 +88,14 @@ struct Session final : public transcribe_session { std::vector ada_scale; // per-layer views int ada_num_delay = -1; - bool encoder_use_flash = false; bool decoder_use_flash = true; - // ---- Incremental streaming state (real reference StaticCache mechanism) -- + // ---- Incremental streaming state (reference StaticCache mechanism) ------- // The encoder transformer runs INCREMENTALLY: new embedder frames attend to // an encoder KV cache (StaticCache, sliding_window=750) under a windowed // mask, each frame processed exactly once. The decoder prefills the prompt - // once then steps one token per audio frame against its persistent KV. The - // cheap conv stem + mel are recomputed over the accumulated buffer and - // sliced (bit-exact via causal conv + frame-independent global mel); a conv - // padding cache / streaming mel is a documented constant-cost follow-up. + // once then steps one token per audio frame against its persistent KV. std::vector stream_pcm; // accumulated raw PCM (no pad) int stream_num_delay = -1; // active num_delay_tokens int stream_min_decode_ms = -1; // partial-decode throttle @@ -111,10 +105,9 @@ 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) - // Conv-stem padding cache (reference VoxtralRealtimeConv1dCacheLayer): feed only - // new mel frames through the conv stem, carrying the conv left-context instead of - // recomputing the whole buffer. conv0 (k3 s1) ← last 2 mel frames; conv1 (k3 s2) - // ← last 1 conv0-output frame. Zeros on the first chunk == whole-buffer left-pad. + // 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. std::vector stream_conv0_cache; // [n_mels, 2] (mel-major, 2 frames) std::vector stream_conv1_cache; // [d_model, 1] (last conv0-out frame) int stream_n_mel_committed = 0; // mel frames fed through the conv stem diff --git a/src/arch/voxtral_realtime/weights.h b/src/arch/voxtral_realtime/weights.h index 468d16cb..838b6022 100644 --- a/src/arch/voxtral_realtime/weights.h +++ b/src/arch/voxtral_realtime/weights.h @@ -1,7 +1,7 @@ // arch/voxtral_realtime/weights.h - Voxtral Realtime (2602) catalog + hparams. // -// INTERNAL to src/arch/voxtral_realtime/. Streaming audio-LLM, distinct -// from the 2507 `voxtral` family. Three-sided weight layout: +// INTERNAL to src/arch/voxtral_realtime/. Streaming audio-LLM, distinct from the +// 2507 `voxtral` family. Three-sided weight layout: // // enc.* causal RoPE audio encoder (2x left-pad causal Conv1d stem + // 32 pre-norm RMSNorm transformer layers; NEOX RoPE theta 1e6 diff --git a/src/arch/whisper/bin_load.cpp b/src/arch/whisper/bin_load.cpp index 30c90bd8..bf44ccad 100644 --- a/src/arch/whisper/bin_load.cpp +++ b/src/arch/whisper/bin_load.cpp @@ -1,10 +1,7 @@ -// arch/whisper/bin_load.cpp - legacy whisper.cpp .bin adapter. -// -// See bin_load.h for the contract. This file synthesizes everything -// the canonical Whisper runtime needs from a parsed WhisperBinModel: -// hparams, capabilities, language list, special token ids, tokenizer, -// frontend buffers, and a ctx_meta populated with canonical-named -// ggml_tensors that the byte-streaming step then fills. +// arch/whisper/bin_load.cpp - legacy whisper.cpp .bin adapter (see bin_load.h +// for the contract). Synthesizes hparams, capabilities, language list, special +// token ids, tokenizer, and frontend buffers from a parsed WhisperBinModel, +// then populates a canonical-named ctx_meta the byte-streaming step fills. #include "bin_load.h" @@ -40,10 +37,8 @@ namespace { constexpr const char * kTag = "whisper-bin"; -// Cosmetic variant string from .bin geometry. Keeps parity with the -// stt.variant a GGUF would carry (whisper-tiny, whisper-medium, etc.) -// for diagnostic logging. Doesn't gate any runtime behavior — that's -// driven by hparams + caps. +// Cosmetic variant string from .bin geometry, for diagnostic logging only +// (parity with the stt.variant a GGUF would carry). Doesn't gate runtime. std::string detect_variant(const transcribe::bin_loader::WhisperBinModel & bm) { const auto & h = bm.hp; std::string base; @@ -69,13 +64,8 @@ std::string detect_variant(const transcribe::bin_loader::WhisperBinModel & bm) { return base; } -// --------------------------------------------------------------------------- -// Canonical Whisper language table (lifted verbatim from whisper.cpp: -// src/whisper.cpp:280-381). Index = whisper language id; for -// is_multilingual variants we install the first `num_languages` -// entries into general.languages-equivalent state. -// --------------------------------------------------------------------------- - +// Canonical Whisper language table (verbatim from whisper.cpp). Index = +// whisper language id; multilingual variants install the first num_languages. const char * const k_whisper_lang_codes[] = { "en","zh","de","es","ru","ko","fr","ja","pt","tr", "pl","ca","nl","ar","sv","it","id","hi","fi","vi", @@ -91,15 +81,9 @@ const char * const k_whisper_lang_codes[] = { constexpr int k_whisper_n_lang_codes = static_cast(sizeof(k_whisper_lang_codes) / sizeof(k_whisper_lang_codes[0])); -// --------------------------------------------------------------------------- -// Special token computation. Mirrors whisper.cpp:1625-1639 verbatim: -// the static defaults below are the multilingual-base layout, then -// for is_multilingual the eot/sot offsets shift by +1 and the -// downstream tokens shift by `num_languages - 98`. .en variants use -// the static defaults unchanged (they sit one slot lower because no -// language tokens occupy the band between sot and translate). -// --------------------------------------------------------------------------- - +// Special token computation (mirrors whisper.cpp). Defaults are the .en / +// multilingual-base layout; for is_multilingual eot/sot shift +1 and the +// downstream tokens shift by num_languages - 98. struct WhisperSpecials { int eot = 50256; int sot = 50257; @@ -131,19 +115,10 @@ WhisperSpecials compute_specials(int n_vocab) { return s; } -// --------------------------------------------------------------------------- -// Suppression token list. The HF generation_config.suppress_tokens -// for whisper is a fixed set of ~88 token ids that should never be -// emitted in transcripts (special tokens, stray punctuation IDs). -// We hardcode it here so .bin-loaded models match transcript quality -// of GGUFs converted via convert-whisper.py without us having to -// re-derive it from the .bin (it isn't stored there). -// -// Source: HF generation_config.json for openai/whisper-tiny (the list -// is identical across multilingual variants; only the trailing -// special-token ids would differ for .en, where no <|translate|> / -// <|notimestamps|> exist, but suppressing already-absent ids is a -// no-op). +// HF generation_config.suppress_tokens for whisper (~88 ids never emitted in +// transcripts). Hardcoded here because the .bin doesn't store it; source is HF +// generation_config.json (identical across variants; suppressing ids absent in +// the smaller .en vocab is a no-op). const int32_t k_whisper_suppress_tokens[] = { 1,2,7,8,9,10,14,25,26,27,28,29,31,58,59,60,61,62,63,90,91,92,93, 359,503,522,542,873,893,902,918,922,931,1350,1853,1982,2460, @@ -156,15 +131,9 @@ const int32_t k_whisper_suppress_tokens[] = { constexpr int k_whisper_n_suppress = static_cast(sizeof(k_whisper_suppress_tokens) / sizeof(k_whisper_suppress_tokens[0])); -// --------------------------------------------------------------------------- -// Tensor rename rules. Each entry maps a legacy whisper.cpp tensor -// name to the canonical transcribe.cpp name. The optional -// `squeeze_trailing` flag drops the size-1 trailing dim that -// whisper.cpp uses for conv biases (`[d_model, 1]` → `[d_model]`). -// Layered tensors (encoder.blocks.{i}.* / decoder.blocks.{i}.*) are -// expanded against n_audio_layer / n_text_layer at runtime. -// --------------------------------------------------------------------------- - +// Tensor rename rules: legacy whisper.cpp name -> canonical transcribe.cpp +// name. collapse_to_1d drops the size-1 dim whisper.cpp uses for conv biases. +// Layered block tensors are expanded against n_audio_layer / n_text_layer. struct StaticRename { const char * legacy; const char * canonical; @@ -196,8 +165,7 @@ const LayeredRename k_enc_block_renames[] = { {"encoder.blocks.%d.attn_ln.bias", "enc.blocks.%d.norm_attn.bias" }, {"encoder.blocks.%d.attn.query.weight", "enc.blocks.%d.attn.q.weight" }, {"encoder.blocks.%d.attn.query.bias", "enc.blocks.%d.attn.q.bias" }, - {"encoder.blocks.%d.attn.key.weight", "enc.blocks.%d.attn.k.weight" }, - // No key bias on either side of attention. + {"encoder.blocks.%d.attn.key.weight", "enc.blocks.%d.attn.k.weight" }, // no bias {"encoder.blocks.%d.attn.value.weight", "enc.blocks.%d.attn.v.weight" }, {"encoder.blocks.%d.attn.value.bias", "enc.blocks.%d.attn.v.bias" }, {"encoder.blocks.%d.attn.out.weight", "enc.blocks.%d.attn.out.weight" }, @@ -243,16 +211,9 @@ std::string fmt1(const char * fmt, int idx) { return std::string(buf); } -// --------------------------------------------------------------------------- -// Compute periodic Hann window of length n_fft. Matches the -// "hann_periodic" flavor MelFrontend expects: w[k] = 0.5 * (1 - cos( -// 2*pi*k / N)) for k = 0..N-1 (no division by N-1). -// -// .bin files don't carry the window — whisper.cpp computes it at -// load. We do the same here so install_mel_from_buffers can run with -// both buffers populated. -// --------------------------------------------------------------------------- - +// Periodic Hann window of length n_fft: w[k] = 0.5*(1 - cos(2*pi*k/N)), +// k=0..N-1 (no division by N-1). .bin files don't carry the window, so compute +// it here (as whisper.cpp does) for install_mel_from_buffers. std::vector hann_periodic(int n_fft) { std::vector w(static_cast(n_fft)); const double two_pi_over_n = 2.0 * M_PI / static_cast(n_fft); @@ -263,12 +224,8 @@ std::vector hann_periodic(int n_fft) { return w; } -// --------------------------------------------------------------------------- -// Synthesize a WhisperHParams from the parsed bin model. Every field -// on WhisperHParams that the GGUF path reads from KVs gets a -// hardcoded value here matching the shipped Whisper variants. -// --------------------------------------------------------------------------- - +// Synthesize a WhisperHParams from the parsed bin model. Every field the GGUF +// path reads from KVs gets a hardcoded value matching the shipped variants. void fill_whisper_hparams(const transcribe::bin_loader::WhisperBinModel & bm, WhisperHParams & hp) { @@ -337,17 +294,11 @@ void fill_whisper_hparams(const transcribe::bin_loader::WhisperBinModel & bm, hp.cap_timestamps = true; } -// --------------------------------------------------------------------------- -// Resolve every canonical tensor slot from the parsed bin model: -// looks up each legacy name in `bm.tensors`, creates a ggml_tensor in -// `ctx_meta` with the canonical name + canonical (post-squeeze) ne + -// the source's ggml_type, and records (target, src_offset, src_nbytes) -// into `out_slots` for the byte-streaming step. -// -// Validates that every required canonical slot has a source tensor -// and that no source tensor is left unconsumed. -// --------------------------------------------------------------------------- - +// Resolve every canonical tensor slot from the parsed bin model: look up each +// legacy name, create a ggml_tensor in ctx_meta with the canonical name + ne + +// source type, and record (target, src_offset, src_nbytes) into out_slots for +// the byte-streaming step. Validates that every slot is filled and no source +// tensor is left unconsumed. ggml_tensor * create_canonical_tensor( ggml_context * ctx_meta, const std::string & canonical_name, @@ -360,12 +311,8 @@ ggml_tensor * create_canonical_tensor( for (int i = 0; i < n_dims; ++i) ne[i] = src.ne[i]; if (collapse_to_1d) { - // Whisper.cpp stores conv biases as numpy shape [d_model, 1], - // which the convert script writes reversed to file dims - // [1, d_model] → ne [1, d_model, 1, 1]. The byte payload is - // still just d_model floats laid out contiguously, so we - // collapse to canonical 1D ne [d_model] by finding the single - // dim with size > 1 and using it as ne[0]. + // whisper.cpp conv biases arrive as ne [1, d_model, 1, 1]; the byte + // payload is d_model contiguous floats, so collapse to 1D ne [d_model]. int64_t total = 1; for (int i = 0; i < n_dims; ++i) total *= ne[i]; for (int i = 0; i < GGML_MAX_DIMS; ++i) ne[i] = 1; @@ -405,14 +352,10 @@ transcribe_status resolve_tensors( std::unordered_set consumed; consumed.reserve(by_name.size() * 2); - // The canonical Whisper weights catalog includes two frontend - // tensor slots (frontend.mel_filterbank, frontend.window) that - // aren't carried as tensors in legacy `.bin` files — the - // filterbank is a separate header field, the window is computed - // at runtime. We still need to satisfy build_whisper_weights' - // lookup, so we create empty F32 tensors with the canonical - // shapes here. Their bytes are populated post-backend-alloc by - // populate_frontend_tensors() below. + // frontend.mel_filterbank / frontend.window aren't .bin tensors (the + // filterbank is a header field, the window is computed at load). Create + // empty F32 slots to satisfy build_whisper_weights; bytes are filled + // post-backend-alloc below. { const int64_t n_fft_half = static_cast(hp.fe_n_fft / 2 + 1); @@ -473,11 +416,9 @@ transcribe_status resolve_tensors( } } - // Stray-tensor check: every entry in the .bin should have been - // consumed by some canonical slot. A stray entry usually means the - // .bin carries an extension (e.g. tinydiarize's solm head) that - // our runtime does not yet support — surface it as a hard error - // rather than silently dropping load-bearing weights. + // Stray-tensor check: every .bin entry must be consumed. A stray entry + // usually means an unsupported extension (e.g. tinydiarize's solm head); + // hard-error rather than silently drop load-bearing weights. if (consumed.size() != bm.tensors.size()) { for (const auto & e : bm.tensors) { if (consumed.find(e.name) == consumed.end()) { @@ -493,18 +434,10 @@ transcribe_status resolve_tensors( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- -// Capabilities + language list synthesis. Mirrors what -// apply_family_invariants + read_capability_kv + read_languages_kv -// do for the GGUF path — except we don't have GGUF KVs to read from, -// so we set them from inferred values. -// -// Storage: caps.languages must point at memory whose lifetime matches -// the WhisperModel. We back it with two parallel std::vectors on the -// model: lang_codes (string storage, owned) and a pointer table held -// in a separate vector that caps.languages references. -// --------------------------------------------------------------------------- - +// Capabilities + language list synthesis. Mirrors apply_family_invariants + +// read_capability_kv + read_languages_kv for the GGUF path, but from inferred +// values since there are no GGUF KVs. lang_codes (owned on the model) backs the +// caps.languages pointer table; its lifetime matches the WhisperModel. void apply_caps_and_languages(WhisperModel & m, const transcribe::bin_loader::WhisperBinModel & bm) { @@ -528,12 +461,9 @@ void apply_caps_and_languages(WhisperModel & m, } else { m.lang_codes.emplace_back("en"); } - // PR2: lang_token_ids stays empty for non-multilingual models; - // for multilingual, populate by ID arithmetic. Whisper packs - // language tokens contiguously after <|sot|>: lang_id i has - // token id sot + 1 + i. The .bin vocab does NOT carry the - // "<|en|>" strings — whisper.cpp synthesizes them at load — so - // we cannot use tokenizer.find() and must compute the ids. + // lang_token_ids by ID arithmetic: whisper packs language tokens + // contiguously after <|sot|> (lang_id i -> sot + 1 + i). The .bin vocab + // doesn't carry the "<|en|>" strings, so tokenizer.find() can't be used. m.lang_token_ids.clear(); if (bm.is_multilingual) { const WhisperSpecials sp = compute_specials(bm.hp.n_vocab); @@ -544,32 +474,19 @@ void apply_caps_and_languages(WhisperModel & m, } } - // Install lang_codes into the public capability surface via the - // shared model::set_languages() helper. This makes - // transcribe_get_model_capabilities()->languages / - // ->n_languages observable for .bin-loaded models, matching what - // GGUF whisper produces from general.languages. + // Publish to the capability surface so transcribe_get_model_capabilities() + // ->languages/->n_languages matches the GGUF path. m.set_languages(m.lang_codes); } -// --------------------------------------------------------------------------- -// Build the Tokenizer via the in-memory raw-bytes loader. The .bin -// vocab is tiktoken-style (raw UTF-8 byte sequences, no GPT-2 byte- -// to-unicode remap), so the raw-bytes decode mode is the right -// strategy. has_encoder() returns true: the tiktoken-style encoder -// uses piece_to_id_ as its rank table directly. +// Build the Tokenizer via the raw-bytes loader (the .bin vocab is +// tiktoken-style: raw UTF-8, no GPT-2 byte-to-unicode remap). // -// We then synthesize the canonical Whisper special-token literals -// ("<|en|>", "<|notimestamps|>", "<|0.00|>", …) into the tokenizer's -// special-piece map so find() can resolve them. The .bin file itself -// only stores the ~50256 base vocab pieces; whisper.cpp synthesizes -// the special strings at load time too. Without these entries the -// run path's initial_prompt rejection check would miss literals like -// "<|en|>" embedded in user text — leading to the literal getting -// byte-encoded as raw context noise instead of being rejected -// symmetrically with the GGUF path. -// --------------------------------------------------------------------------- - +// The .bin stores only the ~50256 base pieces, so synthesize the canonical +// special-token literals ("<|en|>", "<|notimestamps|>", "<|0.00|>", ...) into +// the special-piece map (as whisper.cpp does). Without them the run path's +// initial_prompt rejection check would miss such literals in user text and +// byte-encode them as context noise instead of rejecting symmetrically. transcribe_status install_tokenizer(WhisperModel & m, const transcribe::bin_loader::WhisperBinModel & bm) { @@ -600,10 +517,8 @@ transcribe_status install_tokenizer(WhisperModel & } } - // 1501 timestamp tokens at <|0.00|>, <|0.02|>, …, <|30.00|>. - // ids beg + 0 .. beg + 1500. Format matches HF Whisper's - // tokenizer added_tokens convention exactly so a literal - // "<|0.00|>" or "<|30.00|>" in user text is detected. + // 1501 timestamp tokens <|0.00|>..<|30.00|> at ids beg+0..beg+1500. Format + // matches HF's added_tokens convention so such literals are detected. char tsbuf[16]; for (int i = 0; i <= 1500; ++i) { std::snprintf(tsbuf, sizeof(tsbuf), "<|%.2f|>", @@ -713,18 +628,14 @@ transcribe_status load_from_bin(const char * path, return st; } - // ---- Frontend: populate vestigial ggml tensors and install - // MelFrontend with the buffers. ---- + // ---- Frontend: fill the catalog tensors and install MelFrontend ---- { std::vector filterbank = bm.mel_filterbank; // copy std::vector window = hann_periodic(m->hparams.fe_n_fft); - // Vestigial: weights.frontend.* slots are populated for catalog - // completeness. The runtime mel pipeline reads from m->mel. - // The parser already enforces n_fft_filters == 201 and - // n_mel_filters == hp.n_mels, so the byte counts match by - // construction; the asserts here are defense in depth so a - // future parser regression can't turn into a backend OOB write. + // weights.frontend.* are populated for catalog completeness only (the + // runtime mel pipeline reads from m->mel). The size checks are defense + // in depth against a future parser regression causing a backend OOB. if (m->weights.frontend.mel_filterbank != nullptr) { const size_t want = ggml_nbytes(m->weights.frontend.mel_filterbank); const size_t have = filterbank.size() * sizeof(float); diff --git a/src/arch/whisper/bin_load.h b/src/arch/whisper/bin_load.h index 0d175005..5ddde41b 100644 --- a/src/arch/whisper/bin_load.h +++ b/src/arch/whisper/bin_load.h @@ -1,23 +1,12 @@ -// arch/whisper/bin_load.h - legacy whisper.cpp .bin adapter for the -// Whisper architecture. +// arch/whisper/bin_load.h - legacy whisper.cpp .bin adapter. INTERNAL, +// consumed by transcribe.cpp (magic-byte dispatch) and bin_load.cpp. // -// INTERNAL. Consumed by transcribe.cpp (magic-byte dispatch) and -// arch/whisper/bin_load.cpp. -// -// The .bin parser (transcribe-bin-loader.{h,cpp}) is architecture- -// neutral: it knows the byte layout but nothing about WhisperHParams, -// WhisperWeights, or the canonical tensor names this codebase uses. +// The .bin parser (transcribe-bin-loader.{h,cpp}) is architecture-neutral. // This module is the whisper-specific bridge: it consumes a parsed -// WhisperBinModel, synthesizes hparams + capabilities + special-token -// ids + language list, populates the Tokenizer via the decode-only -// raw-bytes path, builds ctx_meta with canonical-named tensors, allocs -// backend memory, and streams tensor bytes from the .bin into the -// allocated slots. -// -// Once load_from_bin returns, the resulting transcribe_model* is -// indistinguishable from one produced by the GGUF whisper_load path -// — same WhisperModel struct, same encoder/decoder runtime, same -// public ABI. +// WhisperBinModel, synthesizes hparams/capabilities/special-tokens/languages, +// populates the Tokenizer, builds ctx_meta with canonical-named tensors, and +// streams tensor bytes into the allocated slots. The resulting +// transcribe_model* is indistinguishable from the GGUF whisper_load path. #pragma once diff --git a/src/arch/whisper/capabilities.cpp b/src/arch/whisper/capabilities.cpp index ca6e92ed..478df099 100644 --- a/src/arch/whisper/capabilities.cpp +++ b/src/arch/whisper/capabilities.cpp @@ -9,35 +9,19 @@ void apply_family_invariants(transcribe_model & model) { caps.native_sample_rate = 16000; - // Multilingual whisper variants support both lang-detect and the - // <|translate|> task token. The shared read_capability_kv() call - // in the load path will override these from GGUF if the converter - // emitted stt.capability.{lang_detect,translate}; the defaults - // here are the correct answer for the multilingual variants we - // ship first. The .en variants will carry - // stt.capability.{lang_detect,translate}=false and overwrite. + // Multilingual defaults; read_capability_kv() in the load path overrides + // these from GGUF (.en variants carry lang_detect/translate=false). caps.supports_language_detect = true; caps.supports_translate = true; - // supports_streaming stays at its zero-init default (false); whisper - // has no streaming hooks. Left implicit to match the other - // non-streaming families rather than restating the default here. - - // Segment timestamps via the timestamp-token stream (<|t=0.00|> … - // <|t=30.00|>, ids 50364+). Word-level timestamps via DTW over - // cross-attention alignment heads are not yet implemented — the - // model is capable, but the DTW path is a Stage-7 follow-up. Cap - // at SEGMENT so the dispatcher rejects premature requests for - // WORD-grain timing rather than silently falling back. + // supports_streaming left at its zero-init default (false). + + // Segment timestamps via the timestamp-token stream (ids 50364+). + // Word-level DTW timing is not yet implemented; cap at SEGMENT so the + // dispatcher rejects WORD-grain requests rather than falling back. caps.max_timestamp_kind = TRANSCRIBE_TIMESTAMPS_SEGMENT; - // Whisper run knobs are reached through transcribe_whisper_run_ext - // (via transcribe_run_params::family). Stage 2 lights up temperature - // fallback + long-form decoding + per-chunk decoding trace. Stage 1 - // already shipped cancellation (abort callback between chunks and - // between decode steps). Stage 3 wires initial_prompt (text or - // pre-tokenized) and condition_on_prev_tokens (cross-chunk - // coherence under <|startofprev|>). No PNC/ITN runtime toggle — - // whisper emits whatever its training distribution produces. + // Run knobs are reached through transcribe_whisper_run_ext. No PNC/ITN + // runtime toggle — whisper emits whatever its training distribution does. transcribe::set_feature(&model, TRANSCRIBE_FEATURE_INITIAL_PROMPT, true); transcribe::set_feature(&model, TRANSCRIBE_FEATURE_TEMPERATURE_FALLBACK, true); transcribe::set_feature(&model, TRANSCRIBE_FEATURE_LONG_FORM, true); diff --git a/src/arch/whisper/decoder.cpp b/src/arch/whisper/decoder.cpp index 33c7dbe7..e0a96f7c 100644 --- a/src/arch/whisper/decoder.cpp +++ b/src/arch/whisper/decoder.cpp @@ -1,4 +1,4 @@ -// arch/whisper/decoder.cpp - Whisper decoder graph builder (prefill). +// arch/whisper/decoder.cpp - Whisper decoder graph builder. // // Mirrors the transformers reference (WhisperDecoderLayer): // @@ -11,12 +11,8 @@ // logits_raw = proj_out(x) # tied: W = embed_tokens.weight // logits = log_softmax(logits_raw) // -// Whisper has NO post-embed LayerNorm (cohere does — that is the -// notable shape difference vs. cohere/decoder.cpp). -// -// q/v/out carry bias on both self and cross attention; k_proj has NO -// bias. The shared `mha_decoder` helper takes nullable bias slots and -// skips the add when null, matching encoder.cpp's `mha_encoder`. +// NO post-embed LayerNorm (unlike cohere). q/v/out carry bias, k does NOT +// (self and cross); the mha_* helpers take nullable bias slots. #include "decoder.h" @@ -72,8 +68,7 @@ ggml_tensor * mha_decoder(ggml_context * ctx, ggml_tensor * q = ggml_mul_mat(ctx, q_w, x); if (q_b != nullptr) q = ggml_add(ctx, q, q_b); - ggml_tensor * k = ggml_mul_mat(ctx, k_w, source); - // k has no bias. + ggml_tensor * k = ggml_mul_mat(ctx, k_w, source); // k has no bias ggml_tensor * v = ggml_mul_mat(ctx, v_w, source); if (v_b != nullptr) v = ggml_add(ctx, v, v_b); @@ -171,35 +166,23 @@ ggml_tensor * build_block(ggml_context * ctx, return x; } -// --------------------------------------------------------------------------- -// KV-cached helpers -// --------------------------------------------------------------------------- +// KV-cached helpers. // // Cache layout (flat 1D tensors, one slot per layer, per role): -// // self_k / self_v : [hidden * n_layer * n_ctx] // layer il offset: il * n_ctx * hidden -// position n slot within layer: -// il * n_ctx * hidden + n * hidden -// +// position n slot: il * n_ctx * hidden + n * hidden // cross_k / cross_v : [hidden * n_layer * T_enc] // layer il offset: il * T_enc * hidden -// -// Shape follows cohere's layout exactly so the ggml patterns transfer -// unchanged; whisper's only quirk is k_proj has no bias on both self -// and cross (cohere has k_b which may also be null). // Self-attention reading/writing the self KV cache. // // x: [d_model, n_tokens] input for new tokens only // mask: [n_kv, n_tokens] f16 additive mask, or nullptr -// il: layer index (for cache offset) -// n_past: number of tokens already in cache -// n_tokens: number of new tokens being processed +// n_past: tokens already in cache; n_tokens: new tokens this call // -// Writes new K/V into the cache at [il, n_past..n_past+n_tokens), -// then reads the full [il, 0..n_past+n_tokens) window for attention. -// Returns [d_model, n_tokens]. +// Writes new K/V into the cache at [il, n_past..n_past+n_tokens), then reads +// the full [il, 0..n_past+n_tokens) window. Returns [d_model, n_tokens]. ggml_tensor * mha_self_cached(ggml_context * ctx, ggml_cgraph * gf, ggml_tensor * x, @@ -220,17 +203,15 @@ ggml_tensor * mha_self_cached(ggml_context * ctx, const int head_dim = d_model / n_heads; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); const int n_ctx = kv_cache.n_ctx; - // n_kv may be padded to a FA-friendly multiple beyond n_past + - // n_tokens. Trailing slots are masked to -inf via the caller's - // mask so they do not contribute. The cache write below still - // touches only the [n_past, n_past+n_tokens) range. + // n_kv may be padded beyond n_past + n_tokens (FA-friendly multiple); + // trailing slots are masked to -inf. The cache write still touches only + // [n_past, n_past+n_tokens). - // Q / K / V projections for the *new* tokens only. + // Q / K / V projections for the new tokens only. ggml_tensor * Qcur = ggml_mul_mat(ctx, q_w, x); if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b); ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x); - // k has no bias on whisper. ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x); if (v_b != nullptr) Vcur = ggml_add(ctx, Vcur, v_b); @@ -330,7 +311,6 @@ ggml_tensor * mha_self_step(ggml_context * ctx, if (q_b != nullptr) Qcur = ggml_add(ctx, Qcur, q_b); ggml_tensor * Kcur = ggml_mul_mat(ctx, k_w, x); - // k has no bias on whisper. ggml_tensor * Vcur = ggml_mul_mat(ctx, v_w, x); if (v_b != nullptr) Vcur = ggml_add(ctx, Vcur, v_b); @@ -405,12 +385,7 @@ ggml_tensor * mha_self_step(ggml_context * ctx, } // Cross-attention reading the pre-populated cross KV cache. -// -// x: [d_model, n_tokens] query input -// il: layer index -// T_enc: number of encoder frames -// -// Returns: [d_model, n_tokens]. +// x: [d_model, n_tokens] query input -> returns [d_model, n_tokens]. ggml_tensor * mha_cross_cached(ggml_context * ctx, ggml_tensor * x, WhisperKvCache & kv_cache, @@ -427,11 +402,9 @@ ggml_tensor * mha_cross_cached(ggml_context * ctx, const float scale = 1.0f / std::sqrt(static_cast(head_dim)); const int64_t n_tokens = x->ne[1]; - // Cross-attn cache is laid out at T_enc_pad rows per layer; the - // trailing T_enc_pad - T_enc rows hold zeros (cleared at init, - // never written by build_cross_kv_graph). FA reads the padded - // shape — the unmasked padded slots contribute a small dilution - // but produce a kernel-friendly sequence dim. + // Cross cache is T_enc_pad rows per layer; trailing rows hold zeros (never + // written). FA reads the padded shape — kernel-friendly seq dim at the cost + // of small dilution unless the caller masks the padded slots. const int T_enc_pad = kv_cache.T_enc_pad > 0 ? kv_cache.T_enc_pad : T_enc; (void)T_enc; @@ -519,8 +492,8 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, named(db.encoder_out_in, "dec.encoder_out"); ggml_set_input(db.encoder_out_in); - // Causal mask declared as F32 input; cast to F16 inside the graph - // because flash_attn_ext / soft_max_ext both want f16. + // Causal mask declared F32, cast to F16 inside (flash_attn_ext / + // soft_max_ext want f16). ggml_tensor * causal_mask = nullptr; if (seq_len > 1) { db.causal_mask_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, @@ -530,20 +503,17 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, causal_mask = ggml_cast(ctx, db.causal_mask_in, GGML_TYPE_F16); } - // ---- Token embedding -------------------------------------------- - // token_embd_w is [d_model, vocab]; ggml_get_rows with i32 indices - // [seq_len] yields [d_model, seq_len]. + // ---- Token embedding ---- + // get_rows(token_embd_w [d_model, vocab], ids [seq_len]) -> [d_model, seq_len]. ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w, db.token_ids_in); named(tok_emb, "dec.token_emb"); transcribe::debug::mark_tensor_for_dump(tok_emb); db.dumps.token_emb = tok_emb; - // ---- Positional embedding --------------------------------------- - // pos_emb_w is [d_model, max_target_positions]. The transformers - // reference takes weight[past:past+seq_len]; for prefill past=0 so - // a 2D view of the leading seq_len rows is exactly the same data - // (and avoids needing an i32 input for position ids). + // ---- Positional embedding ---- + // Reference takes pos_emb_w[past:past+seq_len]; prefill is past=0, so a 2D + // view of the leading seq_len rows is the same data (no pos-id input). ggml_tensor * pos_emb = w.dec_top.pos_emb_w; if (seq_len != hp.dec_max_target_positions) { pos_emb = ggml_view_2d(ctx, w.dec_top.pos_emb_w, @@ -554,13 +524,13 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, transcribe::debug::mark_tensor_for_dump(pos_emb); db.dumps.pos_emb = pos_emb; - // ---- Embedding sum (no LayerNorm) ------------------------------- + // ---- Embedding sum (no LayerNorm) ---- ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb); named(x, "dec.embed_sum"); transcribe::debug::mark_tensor_for_dump(x); db.dumps.embed_sum = x; - // ---- Decoder blocks --------------------------------------------- + // ---- Decoder blocks ---- const int n_blocks = static_cast(w.dec_blocks.size()); db.dumps.block_outs.reserve(static_cast(n_blocks)); for (int i = 0; i < n_blocks; ++i) { @@ -574,15 +544,14 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, db.dumps.block_outs.push_back(x); } - // ---- Final LayerNorm -------------------------------------------- + // ---- Final LayerNorm ---- x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b); named(x, "dec.out_before_head"); transcribe::debug::mark_tensor_for_dump(x); db.dumps.out_before_head = x; - // ---- Logits head (tied weight, no bias) ------------------------- - // mul_mat(W, x) where W=[d_model, vocab] and x=[d_model, seq_len] - // gives [vocab, seq_len]. + // ---- Logits head (tied weight, no bias) ---- + // mul_mat(W [d_model, vocab], x [d_model, seq_len]) -> [vocab, seq_len]. ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x); named(logits_raw, "dec.logits_raw"); transcribe::debug::mark_tensor_for_dump(logits_raw); @@ -596,9 +565,7 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, db.out = logits; ggml_set_output(db.out); - // 8192 leaves headroom for 4-block whisper-tiny (and for larger - // variants up to whisper-large-v3 at 32 blocks; each block adds - // ~20 ops with self+cross+ffn). + // 8192 leaves headroom up to large-v3 (32 blocks, ~20 ops each). db.graph = ggml_new_graph_custom(ctx, 8192, false); if (db.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -607,16 +574,12 @@ DecoderBuild build_decoder_prefill_graph(ggml_context * ctx, } ggml_build_forward_expand(db.graph, db.out); - // Suppress unused-variable warning when vocab is only used in the - // bounds check above (kept for documentation in the catalog). (void)vocab; return db; } -// --------------------------------------------------------------------------- // KV-cached path: cross-KV precompute + prompt/step decoder graph. -// --------------------------------------------------------------------------- DecoderBuild build_cross_kv_graph(ggml_context * ctx, const WhisperWeights & w, @@ -638,9 +601,8 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, const int d_model = hp.dec_d_model; - // View the persistent encoder_out tensor into this compute_ctx. - // No data is copied — the view shares storage with the source, - // which the encoder graph populated on the same backend. + // View the persistent encoder_out tensor into this compute_ctx (no copy; + // shares storage with the encoder-populated source on the same backend). db.encoder_out_in = ggml_view_2d(ctx, encoder_out, d_model, T_enc, encoder_out->nb[1], 0); @@ -653,9 +615,8 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, return db; } - // Layer slots in the flat cross cache are spaced by T_enc_pad - // (allocator size) but each layer only writes T_enc rows; the - // trailing T_enc_pad - T_enc rows stay zero from buffer_clear. + // Layer slots are spaced by T_enc_pad (allocator size) but each writes + // only T_enc rows; trailing T_enc_pad - T_enc rows stay zero. const int T_enc_pad = kv_cache.T_enc_pad > 0 ? kv_cache.T_enc_pad : T_enc; @@ -663,7 +624,7 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, for (int il = 0; il < n_layers; ++il) { const auto & blk = w.dec_blocks[il]; - // Whisper: cross k has no bias; cross v has a bias. + // Cross k has no bias; cross v does. ggml_tensor * Kcross = ggml_mul_mat(ctx, blk.cross_k_w, db.encoder_out_in); ggml_tensor * Vcross = ggml_mul_mat(ctx, blk.cross_v_w, @@ -672,10 +633,8 @@ DecoderBuild build_cross_kv_graph(ggml_context * ctx, Vcross = ggml_add(ctx, Vcross, blk.cross_v_b); } - // Write the [d_model, T_enc] projections into the flat cache - // slot for this layer. Offset in elements = il * T_enc_pad * - // d_model so the layer slots in the padded cache do not - // overlap; the size of each write is T_enc * d_model. + // Write [d_model, T_enc] into this layer's slot at element offset + // il * T_enc_pad * d_model (slots don't overlap). const size_t k_elem = ggml_element_size(kv_cache.cross_k); const size_t v_elem = ggml_element_size(kv_cache.cross_v); @@ -726,11 +685,9 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } - // Pad the K/V view length up to a FA-friendly multiple. Trailing - // slots in [n_kv_active, n_kv) hold either zeros (first use of a - // freshly-cleared cache) or stale K/V values from prior tiers / - // chunks; the mask zeroes them out so they do not contribute. - // Clamp to kv_cache.n_ctx so we do not read past the allocation. + // Pad the K/V view length to a FA-friendly multiple. Trailing slots in + // [n_kv_active, n_kv) hold zeros or stale K/V; the mask zeroes them. + // Clamp to n_ctx so we don't read past the allocation. const int kv_pad_eff = kv_pad > 0 ? kv_pad : 1; int n_kv = n_kv_active; if (kv_pad_eff > 1) { @@ -748,26 +705,16 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, named(db.token_ids_in, "dec.token_ids"); ggml_set_input(db.token_ids_in); - // Position ids as an input so the step loop can upload a single - // int32 (the absolute position) without rebuilding the graph - // structure. For prompt pass (n_past=0, n_tokens=4) the caller - // uploads [0, 1, 2, 3]. + // Position ids as an input so the step loop uploads absolute positions + // without rebuilding the graph (prompt pass uploads [0..n_tokens)). ggml_tensor * pos_ids_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); named(pos_ids_in, "dec.pos_ids"); ggml_set_input(pos_ids_in); - // Self-attention mask. Two reasons we need one: - // (1) prompt pass (n_tokens > 1) needs a causal mask - // (2) padded step (kv_pad > 1) needs -inf at trailing slots - // The single mask covers both: shape [n_kv, n_tokens] with -inf - // for k >= n_kv_active (covers padding) and -inf for k > q + - // n_past (covers causality when n_tokens > 1). Caller in model.cpp - // populates the values. - // - // When n_tokens == 1 AND there is no padding, the mask can be - // omitted (single token attends to the full real cache window). - // Preserve that fast path so non-FA / non-Metal paths skip the - // extra input tensor. + // Self-attention mask [n_kv, n_tokens], covering both causality (prompt + // pass, n_tokens > 1) and padding (kv_pad > 1): -inf for k >= n_kv_active + // and for k > q + n_past. Omitted on the unpadded single-token step (it + // attends the full real cache window); model.cpp fills the values. ggml_tensor * causal_mask = nullptr; const bool need_mask = (n_tokens > 1) || has_padding; if (need_mask) { @@ -781,12 +728,9 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, "decoder_kv mask-null branch requires single-token step"); } - // Cross-attention mask. Only needed when the cross cache has - // been allocated with padding (T_enc_pad > T_enc); the trailing - // padded slots hold zero K/V and would otherwise dilute the - // unmasked FA softmax. -inf at those positions zeros the - // contribution and restores numerical equivalence to an - // unpadded view. + // Cross-attention mask, only when the cross cache is padded (T_enc_pad > + // T_enc): the trailing zero-K/V slots would dilute the FA softmax, so -inf + // there restores equivalence to an unpadded view. ggml_tensor * cross_mask = nullptr; const int T_enc_pad = kv_cache.T_enc_pad > 0 ? kv_cache.T_enc_pad : T_enc; @@ -798,7 +742,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, cross_mask = ggml_cast(ctx, db.cross_mask_in, GGML_TYPE_F16); } - // ---- Token embedding ------------------------------------------- + // ---- Token embedding ---- ggml_tensor * tok_emb = ggml_get_rows(ctx, w.dec_top.token_embd_w, db.token_ids_in); named(tok_emb, "dec.token_emb"); @@ -807,7 +751,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.token_emb = tok_emb; } - // ---- Positional embedding -------------------------------------- + // ---- Positional embedding ---- ggml_tensor * pos_emb = ggml_get_rows(ctx, w.dec_top.pos_emb_w, pos_ids_in); named(pos_emb, "dec.pos_emb"); @@ -816,7 +760,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.pos_emb = pos_emb; } - // ---- Embed sum (no LayerNorm) ---------------------------------- + // ---- Embed sum (no LayerNorm) ---- ggml_tensor * x = ggml_add(ctx, tok_emb, pos_emb); named(x, "dec.embed_sum"); if (n_past == 0) { @@ -824,8 +768,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.embed_sum = x; } - // Pre-create the graph so mha_self_cached can use - // ggml_build_forward_expand to wire the cache writes into it. + // Pre-create the graph so mha_self_cached can wire cache writes into it. db.graph = ggml_new_graph_custom(ctx, 8192, false); if (db.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -833,7 +776,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } - // ---- Decoder blocks -------------------------------------------- + // ---- Decoder blocks ---- const int n_blocks = static_cast(w.dec_blocks.size()); db.dumps.block_outs.reserve(static_cast(n_blocks)); for (int i = 0; i < n_blocks; ++i) { @@ -880,7 +823,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, } } - // ---- Final LayerNorm ------------------------------------------- + // ---- Final LayerNorm ---- x = layer_norm(ctx, x, w.dec_top.final_norm_w, w.dec_top.final_norm_b); named(x, "dec.out_before_head"); if (n_past == 0) { @@ -888,7 +831,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.out_before_head = x; } - // ---- Logits head (tied weight, no bias) ------------------------ + // ---- Logits head (tied weight, no bias) ---- ggml_tensor * logits_raw = ggml_mul_mat(ctx, w.dec_top.token_embd_w, x); named(logits_raw, "dec.logits_raw"); if (n_past == 0) { @@ -913,10 +856,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, return db; } -// --------------------------------------------------------------------------- // Static-topology single-token step graph (GPU dispatch path). -// --------------------------------------------------------------------------- - StepBuild build_step_graph(ggml_context * ctx, const WhisperWeights & w, const WhisperHParams & hp, @@ -1031,10 +971,7 @@ StepBuild build_step_graph(ggml_context * ctx, return sb; } -// =========================================================================== -// Offline batched decode (B utterances) -// =========================================================================== - +// Offline batched decode (B utterances). namespace { // Batched self-attention step. x: [d_model, B] (one token per utterance). @@ -1170,7 +1107,7 @@ DecoderBuild build_cross_kv_graph_batched(ggml_context * ctx, for (int il = 0; il < n_layers; ++il) { const auto & blk = w.dec_blocks[il]; - // Whisper: cross k has no bias; cross v has a bias. + // Cross k has no bias; cross v does. ggml_tensor * Kc = ggml_mul_mat(ctx, blk.cross_k_w, db.encoder_out_in); ggml_tensor * Vc = ggml_mul_mat(ctx, blk.cross_v_w, db.encoder_out_in); if (blk.cross_v_b != nullptr) Vc = ggml_add(ctx, Vc, blk.cross_v_b); diff --git a/src/arch/whisper/decoder.h b/src/arch/whisper/decoder.h index 62034b13..cc662b61 100644 --- a/src/arch/whisper/decoder.h +++ b/src/arch/whisper/decoder.h @@ -1,29 +1,19 @@ // arch/whisper/decoder.h - Whisper decoder graph builder. // -// The decoder is an autoregressive Transformer with: -// - Token embedding + learned positional embedding (NO post-embed LayerNorm) -// - N decoder layers (pre-LN self-attn + pre-LN cross-attn + pre-LN FFN) -// - Final LayerNorm -// - Linear head (tied to token embedding, NO bias) + log-softmax -// -// Two graph families are exposed: -// -// 1. build_decoder_prefill_graph() — non-cached prefill, used only -// for language detection (a single SOT-only forward pass on the -// first chunk before the cross-KV cache exists). Cross-K/V are -// recomputed from the encoder output inside the graph. +// Autoregressive Transformer: token embedding + learned positional embedding +// (NO post-embed LayerNorm), N pre-LN layers (self-attn + cross-attn + FFN), +// final LayerNorm, linear head tied to token embedding (NO bias) + log-softmax. // +// Two graph families: +// 1. build_decoder_prefill_graph() — non-cached, used only for language +// detection (SOT-only forward on the first chunk before the cross-KV +// cache exists); cross-K/V recomputed inside the graph. // 2. build_cross_kv_graph() + build_decoder_graph_kv() — KV-cached -// autoregressive path. The production decoder driver. Cross-K/V -// are precomputed once per chunk into the cross cache; self-K/V -// are written on each step into the self cache. The same graph -// handles both the prompt pass (n_tokens > 1, n_past = 0) and -// per-step generation (n_tokens = 1, n_past = current). All -// dec.* dump points are wired on the prompt pass so validate.py -// tolerances apply directly. +// production path. Cross-K/V precomputed once per chunk; self-K/V written +// per step. One graph handles both prompt pass (n_tokens>1, n_past=0) and +// per-step generation (n_tokens=1, n_past=current). // -// Whisper attention quirk (matches encoder.cpp): q/v/out projections -// carry bias; k_proj has NO bias on both self- and cross-attention. +// Attention quirk (as encoder.cpp): q/v/out carry bias, k does NOT. #pragma once @@ -102,10 +92,7 @@ DecoderBuild build_cross_kv_graph(ggml_context * compute_ctx, // prev total). Self-K/V are written into the self cache at position // n_past; the full [0..n_past+n_tokens) window is read back for // attention. Cross-attention reads from the pre-populated cross cache. -// -// On the prompt pass dumps are wired for every dec.* tensor so the -// validate.py harness keeps working. Step passes produce the same -// graph shape minus the dumps. +// The prompt pass wires every dec.* dump; step passes omit them. // // skip_log_softmax=true outputs pre-softmax logits (argmax-invariant, // cheaper readback). skip_log_softmax=false matches the reference dump. @@ -159,17 +146,11 @@ StepBuild build_step_graph(ggml_context * compute_ctx, int T_enc, bool use_flash = true); -// --------------------------------------------------------------------------- -// Offline batched decode (B utterances in lockstep). -// -// Mirrors the single-shot cross-KV + step graphs but threads an utterance -// batch on the trailing flash axis (ne[3]). The cross cache is populated -// once per batch from a host-packed [d_model, T_enc_max, B] encoder tensor; -// the step graph emits RAW logits [vocab, B] (NOT an in-graph argmax) so the -// host can apply whisper's suppress / begin_suppress / timestamp / no_speech -// rules per utterance before sampling. Flash-only (the batched view layout +// Offline batched decode (B utterances in lockstep). Mirrors the single-shot +// cross-KV + step graphs but threads the batch on the trailing flash axis +// (ne[3]). Emits RAW logits [vocab, B] so the host applies whisper's per- +// utterance rules before sampling. Flash-only (the batched view layout // requires flash_attn_ext). -// --------------------------------------------------------------------------- // Build a batched cross-attention K/V precompute graph. encoder_out_in is a // [d_model, T_enc_max, B] f32 input the caller fills with zero-padded diff --git a/src/arch/whisper/encoder.cpp b/src/arch/whisper/encoder.cpp index 51fa6b7d..8fe8ade8 100644 --- a/src/arch/whisper/encoder.cpp +++ b/src/arch/whisper/encoder.cpp @@ -7,24 +7,17 @@ // -> transpose + cont to [T, n_mels] for ggml_conv_1d // -> conv1 (k=3, s=1, p=1) + GELU [T=3000, d_model] // -> conv2 (k=3, s=2, p=1) + GELU [T=1500, d_model] -// -> transpose + cont to [d_model, T=1500] (channel-innermost, -// matching reference -// dump layout and the -// natural transformer -// layout) +// -> transpose + cont to [d_model, T=1500] (channel-innermost) // -> + learned positional embedding [d_model, T=1500] -// -> 4x pre-LN transformer block +// -> Nx pre-LN transformer block // -> final LayerNorm // -> enc.final [d_model, T=1500] // -// The reference dumps pin the layout: enc.mel.in numpy shape is -// (T, n_mels) and every post-stem activation is (T, d_model). In -// ggml ne terms (fast innermost), that's [n_mels, T] for the mel -// input and [d_model, T] for everything after the stem transpose. +// ggml ne (fast innermost): [n_mels, T] for the mel input, [d_model, T] +// for everything after the stem transpose. // -// Whisper attention quirk: q / v / out carry bias; k has NO bias. -// The `mha` helper below takes all four bias slots as nullable -// ggml_tensor *'s and skips the add when the slot is null. +// Whisper attention quirk: q/v/out carry bias; k has NO bias. The mha helper +// takes nullable bias slots and skips the add when null. #include "encoder.h" @@ -62,14 +55,8 @@ ggml_tensor * add_conv1d_bias(ggml_context * ctx, } // Standard multi-head self-attention without relative position. -// Whisper-specific: q/v/out carry bias; k does NOT. The helper -// accepts nullable bias slots so the call site just passes whatever -// weights.cpp populated. -// -// x: [d_model, T] -// mask: additive mask (f16), or nullptr for unmasked / self-attn -// in the encoder (no causal mask — encoder is bidirectional) -// Returns: [d_model, T] +// x: [d_model, T] -> returns [d_model, T]. Encoder is bidirectional +// (no causal mask). ggml_tensor * mha_encoder(ggml_context * ctx, ggml_tensor * x, ggml_tensor * q_w, ggml_tensor * q_b, @@ -84,12 +71,11 @@ ggml_tensor * mha_encoder(ggml_context * ctx, const float scale = 1.0f / std::sqrt(static_cast(head_dim)); const int64_t T = x->ne[1]; - // Q, K, V projections. + // Q, K, V projections (k has no bias). ggml_tensor * q = ggml_mul_mat(ctx, q_w, x); if (q_b != nullptr) q = ggml_add(ctx, q, q_b); ggml_tensor * k = ggml_mul_mat(ctx, k_w, x); - // k has no bias. ggml_tensor * v = ggml_mul_mat(ctx, v_w, x); if (v_b != nullptr) v = ggml_add(ctx, v, v_b); @@ -108,9 +94,7 @@ ggml_tensor * mha_encoder(ggml_context * ctx, ggml_tensor * o; if (use_flash) { - // flash_attn_ext wants its q/k/v in contiguous form. ggml - // inserts internal cont()s when needed but being explicit - // makes the graph easier to reason about. + // flash_attn_ext wants q/k/v contiguous. ggml_tensor * q_c = ggml_cont(ctx, q); ggml_tensor * k_c = ggml_cont(ctx, k); ggml_tensor * v_c = ggml_cont(ctx, v); @@ -139,8 +123,7 @@ ggml_tensor * mha_encoder(ggml_context * ctx, return o; } -// FFN used by every transformer block: pre-LN wrapped outside, inside -// is just `fc2(GELU(fc1(x)))`. fc1 / fc2 both carry bias in whisper. +// FFN: fc2(GELU(fc1(x))); pre-LN wrapped outside. fc1/fc2 carry bias. ggml_tensor * ffn(ggml_context * ctx, ggml_tensor * x, ggml_tensor * fc1_w, ggml_tensor * fc1_b, @@ -224,10 +207,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, return eb; } - // ---- mel input handle --------------------------------------------- - // - // ggml ne=[n_mels, T] matches the reference dump layout - // enc.mel.in: numpy shape (T, n_mels). The caller uploads the mel + // mel input handle. ggml ne=[n_mels, T]; caller uploads the mel // spectrogram row-major in this layout. eb.mel_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_mels, n_mel_frames); if (eb.mel_in == nullptr) return eb; @@ -236,46 +216,34 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.mel_in = eb.mel_in; transcribe::debug::mark_tensor_for_dump(eb.mel_in); - // ---- conv stem ---------------------------------------------------- - // - // ggml_conv_1d wants data ne=[T, Cin]. Transpose mel from - // [n_mels, T] to [T, n_mels], materialize with cont, feed to conv1. + // conv stem. ggml_conv_1d wants data ne=[T, Cin]; transpose mel from + // [n_mels, T] to [T, n_mels], cont, feed to conv1. ggml_tensor * x = ggml_cont(ctx, ggml_transpose(ctx, eb.mel_in)); - // conv1: k=3, stride=1, padding=1 -> [T, d_model] - // Use conf::conv_1d_f32 instead of ggml_conv_1d because whisper - // stores conv kernels in F32; ggml_conv_1d's internal im2col hard- - // codes F16 as the im2col dst dtype and the CPU kernel asserts - // src0->type == F16, crashing at run time. + // conv1: k=3, stride=1, padding=1 -> [T, d_model]. Use conf::conv_1d_f32, + // not ggml_conv_1d: whisper conv kernels are F32, but ggml_conv_1d's + // im2col hard-codes F16 dst and the CPU kernel asserts src0==F16 (crash). x = conf::conv_1d_f32(ctx, w.enc_stem.conv0_w, x, 1, 1, 1); x = add_conv1d_bias(ctx, x, w.enc_stem.conv0_b); - // HF Whisper uses nn.functional.gelu (exact erf form) for both conv - // stems and the FFN; ggml_gelu is the tanh approximation. + // HF Whisper uses exact-erf GELU for conv stems and FFN; ggml_gelu is tanh. x = ggml_gelu_erf(ctx, x); - // Transpose to [d_model, T] so the dump matches reference layout. - x = ggml_cont(ctx, ggml_transpose(ctx, x)); + x = ggml_cont(ctx, ggml_transpose(ctx, x)); // -> [d_model, T] named(x, "enc.conv1.out"); eb.dumps.conv1_out = x; transcribe::debug::mark_tensor_for_dump(x); - // conv2 needs input in [T, Cin] layout again; transpose back. + // conv2: k=3, stride=2, padding=1 -> [T_enc, d_model] (needs [T, Cin]). x = ggml_cont(ctx, ggml_transpose(ctx, x)); - // conv2: k=3, stride=2, padding=1 -> [T_enc, d_model] x = conf::conv_1d_f32(ctx, w.enc_stem.conv1_w, x, 2, 1, 1); x = add_conv1d_bias(ctx, x, w.enc_stem.conv1_b); x = ggml_gelu_erf(ctx, x); - // Back to [d_model, T_enc] for the rest of the encoder. - x = ggml_cont(ctx, ggml_transpose(ctx, x)); + x = ggml_cont(ctx, ggml_transpose(ctx, x)); // -> [d_model, T_enc] named(x, "enc.conv2.out"); eb.dumps.conv2_out = x; transcribe::debug::mark_tensor_for_dump(x); - // ---- positional embedding + add ---------------------------------- - // - // pos_emb weight is stored in GGUF with ne=[d_model, max_source_positions]. - // For T_enc < max_source_positions, take a prefix view. For whisper-tiny - // T_enc is always 1500 = max_source_positions, so this is a no-op view - // in practice. + // positional embedding + add. pos_emb ne=[d_model, max_source_positions]; + // for T_enc < max take a prefix view (no-op when T_enc == max). ggml_tensor * pos_emb = w.enc_top.pos_emb_w; if (T_enc != hp.enc_max_source_positions) { pos_emb = ggml_view_2d(ctx, w.enc_top.pos_emb_w, @@ -291,17 +259,14 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.dumps.embed_out = x; transcribe::debug::mark_tensor_for_dump(x); - // ---- transformer blocks ------------------------------------------ + // transformer blocks const int n_blocks = static_cast(w.enc_blocks.size()); eb.dumps.block_outs.reserve(static_cast(n_blocks)); for (int i = 0; i < n_blocks; ++i) { x = build_block(ctx, x, w.enc_blocks[i], n_heads, d_model, use_flash); - // Name every block so the reference-dump harness can compare - // layer-by-layer. block0 / block_last are stashed separately - // for backwards compatibility with callers that only want the - // spot-check pair; block_outs carries the full vector so a - // caller that wants to dump all N layers can iterate. + // Name every block for layer-by-layer comparison. block0 / block_last + // are stashed separately for callers that only want the spot-check pair. char bname[64]; std::snprintf(bname, sizeof(bname), "enc.block.%d.out", i); named(x, bname); @@ -315,7 +280,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } } - // ---- final layer norm -------------------------------------------- + // final layer norm x = layer_norm(ctx, x, w.enc_top.final_norm_w, w.enc_top.final_norm_b); named(x, "enc.final"); eb.dumps.final_out = x; @@ -324,8 +289,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, eb.out = x; ggml_set_output(eb.out); - // Build the forward cgraph. whisper-tiny is 4 blocks; other variants - // go up to 32 (large). 8192 leaves headroom for the block ops. + // Build the forward cgraph. 8192 leaves headroom up to large (32 blocks). eb.graph = ggml_new_graph_custom(ctx, 8192, false); if (eb.graph == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, diff --git a/src/arch/whisper/encoder.h b/src/arch/whisper/encoder.h index da4ac5c3..acc4119d 100644 --- a/src/arch/whisper/encoder.h +++ b/src/arch/whisper/encoder.h @@ -1,23 +1,12 @@ // arch/whisper/encoder.h - Whisper encoder graph builder. // -// Forward declarations for the encoder graph construction entry point. -// The encoder is a 2-layer Conv1d stem followed by a stack of pre-LN -// transformer blocks and a final LayerNorm; no relative position -// encoding, no convolution module inside the blocks, and no -// enc-dec projection (the encoder output is already in the decoder's -// hidden dim since upstream whisper ties enc_d_model == dec_d_model). +// 2-layer Conv1d stem -> pre-LN transformer blocks -> final LayerNorm; no +// relative position encoding, no in-block conv, no enc-dec projection +// (enc_d_model == dec_d_model upstream). // -// Dumped tensor names follow the Stage-2 reference script -// (dump_reference_whisper_transformers.py): enc.mel.in, enc.conv1.out, -// enc.conv2.out, enc.pos_emb, enc.embed.out, enc.block.{0..N-1}.out -// (emitted for index 0 and the last block only — the first and last -// are the two spot-checks the reference dumps), enc.final. -// -// Layout: all intermediate activations are kept in [d_model, T] ggml -// layout (d_model innermost, T second), matching the -// channel-innermost convention the reference dumps use. The conv stem -// transposes into [T, d_model] for the ggml_conv_1d calls and back -// out afterwards. +// Layout: all intermediate activations are kept in [d_model, T] ggml layout +// (d_model innermost, T second). The conv stem transposes into [T, d_model] +// for the ggml_conv_1d calls and back out afterwards. #pragma once @@ -44,10 +33,8 @@ struct EncoderDumps { ggml_tensor * block_last_out = nullptr; ggml_tensor * final_out = nullptr; - // Every post-block residual stream output, in block order. Populated - // alongside block0_out / block_last_out so the model driver can dump - // every block for validate.py parity against the reference dumps, - // which emit enc.block.{0..N-1}.out for all N layers. + // Every post-block residual stream output, in block order (alongside + // block0_out / block_last_out for callers that want all N layers). std::vector block_outs; }; @@ -68,17 +55,10 @@ struct EncoderBuild { // Build the encoder forward graph in compute_ctx. // -// n_mel_frames is the number of frames in the mel input. Upstream -// whisper always pads-or-trims to 3000; the builder accepts any -// positive even value so unit tests can exercise smaller fixtures. -// -// use_flash selects ggml_flash_attn_ext vs. manual mul_mat + soft_max -// for each transformer block. The caller is responsible for gating -// this against backend support (see TRANSCRIBE_NO_FLASH policy). -// -// backend_name is informational only; no per-backend conv policy is -// applied on whisper (the stem is small, vanilla conv1d is fine on -// every backend). +// n_mel_frames must be positive and even (upstream pads/trims to 3000; smaller +// even values are accepted for test fixtures). use_flash selects +// ggml_flash_attn_ext vs. manual mul_mat + soft_max; the caller gates it +// against backend support. backend_name is informational only. EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const WhisperWeights & weights, const WhisperHParams & hp, diff --git a/src/arch/whisper/model.cpp b/src/arch/whisper/model.cpp index 185dff78..94f7c1ab 100644 --- a/src/arch/whisper/model.cpp +++ b/src/arch/whisper/model.cpp @@ -1,9 +1,8 @@ // arch/whisper/model.cpp - Whisper ASR family handler. // -// Stage-4 Whisper runtime. Numerical validation can still inject the -// reference mel tensor with TRANSCRIBE_WHISPER_MEL_FROM_REF to isolate -// encoder/decoder drift, but normal inference uses the C++ mel frontend -// and the autoregressive decoder path below. +// Normal inference uses the C++ mel frontend and the autoregressive decoder +// path below; TRANSCRIBE_WHISPER_MEL_FROM_REF can inject a reference mel tensor +// to isolate encoder/decoder drift during numerical validation. #include "whisper.h" @@ -279,17 +278,9 @@ namespace { constexpr const char k_default_variant[] = "whisper"; -// Prepare cc->compute_ctx for a graph build of capacity at most -// `mem`. If a context already exists and is at least that large, -// ggml_reset clears the tensor metadata in-place (cheap). Only when -// the existing context is too small do we ggml_free + ggml_init. -// -// Replaces the previous "free + init per graph build" pattern that -// allocated a fresh compute context for the encoder, cross-KV graph, -// every tier's prompt prefill, and every step in the generation -// loop. The free + malloc churn was visible in the per-step build -// timer (~30 us per step on tiny F16) plus allocator pressure on -// very long inputs. +// Prepare cc->compute_ctx for a graph build of capacity at most `mem`. If a +// context already exists and is large enough, ggml_reset clears its metadata +// in-place (cheap); only a too-small context triggers ggml_free + ggml_init. bool ensure_compute_ctx(WhisperSession * cc, size_t mem) { if (cc->compute_ctx != nullptr) { if (cc->compute_ctx_size >= mem) { @@ -422,10 +413,7 @@ bool whisper_perf_enabled() { return s != nullptr && s[0] != '\0' && s[0] != '0'; } -// --------------------------------------------------------------------------- // load -// --------------------------------------------------------------------------- - transcribe_status whisper_load( Loader & loader, const transcribe_model_load_params * params, @@ -447,13 +435,11 @@ transcribe_status whisper_load( return st; } - // Tokenizer. Whisper GGUFs carry a "gpt2"-family tokenizer (HF - // byte-level BPE). Vocab_size stays whatever the GGUF declared — - // DO NOT overwrite from the tokenizer. Whisper's decoder vocab - // (51865 for multilingual) is larger than tokenizer.tokens - // (50258) because language / timestamp / task special tokens are - // emitted by the model but not listed in the tokens array. That - // gap is by design for this family. + // Tokenizer. Whisper GGUFs carry a "gpt2"-family (HF byte-level BPE) + // tokenizer. DO NOT overwrite vocab_size from the tokenizer: the decoder + // vocab (51865 multilingual) exceeds tokenizer.tokens (50258) because + // language/timestamp/task specials are model-emitted but not in the tokens + // array. That gap is by design. if (const transcribe_status st = m->tok.load(loader.gguf()); st != TRANSCRIBE_OK) { @@ -512,7 +498,7 @@ transcribe_status whisper_load( } } - // Stage 2: reopen with no_alloc to build the tensor catalog. + // Reopen with no_alloc to build the tensor catalog. gguf_init_params init_params {}; init_params.no_alloc = true; init_params.ctx = &m->ctx_meta; @@ -582,23 +568,14 @@ transcribe_status whisper_load( return st; } - // Resolve per-language token ids. The decoder's forced-language - // step needs <|{code}|> id per supported language. We hydrate the - // mapping here so a vocab drift surfaces at load, not mid-run. + // Resolve per-language token ids and keep an owned copy of the language + // list (decoder code reads m->lang_codes directly rather than walking the + // capability chain). Hydrating here surfaces a vocab drift at load. // - // We also keep our own owned copy of the language list so callers - // that go through m->lang_codes don't depend on the capabilities - // storage lifetime (set_languages owns the string backing for - // caps.languages[], but accessing it requires walking the capability - // chain — a direct owned vector here is simpler for decoder code). - // - // Non-multilingual (.en) models advertise a single-language list - // ["en"] but their vocab does NOT contain a "<|en|>" token: the - // .en decoder prefix is just <|sot|>, with no language or task - // tokens. So for .en we keep lang_codes (so the run path can - // accept params.language == "en") but leave lang_token_ids empty. - // The run path branches on supports_language_detect to skip the - // <|lang|>/<|task|> emission for .en. + // .en models advertise ["en"] but their vocab has no "<|en|>" token (the + // .en prefix is bare <|sot|>), so keep lang_codes (to accept + // params.language == "en") but leave lang_token_ids empty; the run path + // gates <|lang|>/<|task|> emission on supports_language_detect. m->lang_codes.clear(); m->lang_token_ids.clear(); const bool is_multilingual = m->caps.supports_language_detect; @@ -632,10 +609,7 @@ transcribe_status whisper_load( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // init_context -// --------------------------------------------------------------------------- - transcribe_status whisper_init_context( transcribe_model * model, const transcribe_session_params * params, @@ -650,9 +624,7 @@ transcribe_status whisper_init_context( cc->n_threads = params->n_threads; cc->kv_type = params->kv_type; - // Whisper defaults: both encoder and decoder flash-on. Head dim for - // whisper-tiny is 64 (d_model=384, 6 heads); supported on every - // backend we ship. See WhisperSession for the rationale. + // Whisper defaults: both flash-on (see WhisperSession). Env overrides below. cc->encoder_use_flash = true; cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides( @@ -662,31 +634,22 @@ transcribe_status whisper_init_context( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // run namespace { -// Run the encoder on one mel window and leave the output in the -// backend-resident persistent tensor cc->enc_out.tensor (d_enc * T_enc -// floats, F32). Downstream graphs (cross-KV precompute, KV-cached -// decoder) read from that tensor directly via ggml views — no host -// roundtrip. cc->enc_host is only populated on demand for language -// detection's prefill graph, which still needs a host-side -// encoder_out_in input. Also publishes T_enc on the context. The -// compute_ctx is reset inside this helper; the scheduler is created -// lazily on first call. +// Run the encoder on one mel window, leaving the output in the backend-resident +// cc->enc_out.tensor (F32 [d_enc, T_enc]) that downstream graphs read via views +// (no host roundtrip). cc->enc_host is populated on demand only for language +// detection's prefill graph. Publishes T_enc on the context. // -// mel_data is in encoder layout [n_mels, n_mel_frames] with n_mels -// innermost, matching the ggml ne=[n_mels, n_mel_frames] input tensor. -// For the shipped Whisper variants n_mel_frames is 3000; long-form -// windows shorter than 3000 must be zero-padded on the right before -// calling this helper. +// mel_data is in encoder layout [n_mels, n_mel_frames] (n_mels innermost). +// Shipped variants use n_mel_frames=3000; shorter long-form windows must be +// zero-padded on the right before calling. // -// allow_dumps controls whether encoder intermediates are emitted via -// transcribe::debug. Set true for validate.py runs (one chunk only in -// short-form) and for the first chunk in long-form; false otherwise -// so the long-form loop does not overwrite dumps from the first chunk. +// allow_dumps gates encoder-intermediate emission: true for the first chunk +// (and short-form), false afterwards so the long-form loop doesn't overwrite +// the first chunk's dumps. transcribe_status run_whisper_encoder_on_window( WhisperSession * cc, WhisperModel * cm, @@ -696,9 +659,6 @@ transcribe_status run_whisper_encoder_on_window( bool allow_dumps, int & out_T_enc) { - // Reset per-call compute state. ensure_compute_ctx reuses the - // context across encoder / cross-KV / prompt / step builds — - // ggml_reset between calls instead of free + reinit. const int64_t t_enc_build_start = ggml_time_us(); if (!ensure_compute_ctx(cc, 8 * 1024 * 1024)) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, @@ -713,13 +673,10 @@ transcribe_status run_whisper_encoder_on_window( return TRANSCRIBE_ERR_GGUF; } - // Backend-resident encoder output. Allocate the persistent - // F32 tensor once (or re-allocate if T_enc / d_model changed). - // After build_encoder_graph, append a final ggml_cpy from eb.out - // into a view of cc->enc_out.tensor — both live on the primary - // backend, so this is a single intra-backend memcpy that stays - // off the host. Subsequent graphs (cross-KV, lang-det) read - // from cc->enc_out.tensor directly. + // Backend-resident encoder output: allocate the persistent F32 tensor once + // (re-allocate on T_enc / d_model change), then append a ggml_cpy from + // eb.out into a view of it — a single intra-backend memcpy that stays off + // the host. Subsequent graphs read cc->enc_out.tensor directly. { const int d_enc_g = static_cast(eb.out->ne[0]); const int T_enc_g = static_cast(eb.out->ne[1]); @@ -758,11 +715,10 @@ transcribe_status run_whisper_encoder_on_window( return TRANSCRIBE_ERR_GGUF; } - // Apply the caller's CPU thread count to the scheduler's backends once - // at sched creation — it persists across the reused encoder and decoder - // graphs, and cc->n_threads is fixed for the session. GPU backends ignore - // it; CPU/BLAS backends use it. Without this, whisper's graph compute ran - // at ggml's default thread count regardless of params.n_threads. + // Apply the caller's CPU thread count once at sched creation; it + // persists across the reused encoder/decoder graphs. GPU backends + // ignore it, CPU/BLAS use it. Without this, compute ran at ggml's + // default count regardless of params.n_threads. transcribe::configure_sched_n_threads(cc->sched, cc->n_threads); } ggml_backend_sched_reset(cc->sched); @@ -793,8 +749,8 @@ transcribe_status run_whisper_encoder_on_window( } cc->perf.enc_compute.add(ggml_time_us() - t_enc_compute_start); - // Dumps. Only the first (or only) chunk emits intermediates so - // validate.py against a reference dump directory sees stable output. + // Dumps. Only the first (or only) chunk emits intermediates so the + // reference comparison sees stable output. if (allow_dumps) { auto try_dump = [](const char * name, ggml_tensor * t, const char * stage) @@ -817,11 +773,8 @@ transcribe_status run_whisper_encoder_on_window( try_dump("enc.final", eb.dumps.final_out, "encoder.final"); } - // The persistent enc_out tensor is now populated. Cross-KV - // reads from it directly — no host materialization needed for - // the hot path. Lang-detect (one-shot, first chunk) still uses - // a tensor-set into a graph input; we lazily materialize to - // enc_host there (see chunk loop below). + // enc_out is now populated; cross-KV reads it directly. Lang-detect + // materializes to enc_host lazily (see chunk loop). const int d_enc = static_cast(eb.out->ne[0]); out_T_enc = static_cast(eb.out->ne[1]); cc->enc_T = out_T_enc; @@ -886,9 +839,8 @@ float compute_compression_ratio_hf( // // -INFINITY logits (from suppress/timestamp rules) contribute 0 mass. // -// `scratch` is reused across calls to avoid the per-step double[vocab] -// alloc on the T>0 hot path. Sized lazily; iteration order matches the -// previous code so the multinomial draw is bit-identical. +// `scratch` is reused across calls to avoid the per-step double[vocab] alloc on +// the T>0 hot path; sized lazily. int sample_from_logits( const std::vector & logits, float temperature, @@ -1082,11 +1034,10 @@ transcribe_status load_mel_from_ref(const char * ref_dir, namespace { -// Whisper timestamp logits processor — shared by the serial whisper_run and -// the batched whisper_run_batch so the per-step masking is identical by -// construction. Mirrors transformers' WhisperTimeStampLogitsProcessor; see -// the rule-by-rule rationale at the (former) lambda site. Mutates `logits` -// in place; no-op when want_segment_timestamps is false. +// Whisper timestamp logits processor, shared by whisper_run and +// whisper_run_batch so per-step masking is identical. Mirrors transformers' +// WhisperTimeStampLogitsProcessor. Mutates `logits` in place; no-op when +// want_segment_timestamps is false. void apply_whisper_timestamp_rules( std::vector & logits, const std::vector & generated_ids, @@ -1185,9 +1136,8 @@ struct WhisperSegmentResult { // Port of HF generation_whisper.py:_retrieve_segment, shared by serial + // batched decode. Pure: reads generated_ids + the tokenizer, returns the -// segments/slices/offset instead of mutating session state. See the call -// sites for the rule rationale. (For short-form callers time_offset_ms is 0 -// and only `.segments` is consumed.) +// segments/slices/offset instead of mutating session state. (Short-form +// callers pass time_offset_ms 0 and consume only `.segments`.) WhisperSegmentResult whisper_retrieve_segment( const std::vector & generated_ids, const transcribe::Tokenizer & tok, @@ -1336,46 +1286,28 @@ transcribe_status whisper_run( transcribe::debug::init(); - // Reset per-stage timing counters; the summary printed at end-of-run - // is opt-in via TRANSCRIBE_WHISPER_PROFILE=1 but the counters - // themselves are always populated (the timestamps are negligible). cc->perf.reset(); - // ----- Mel frontend ------------------------------------------------ + // ----- Mel frontend ----- // - // 1. TRANSCRIBE_WHISPER_MEL_FROM_REF= — debug-only knob. - // Reads the reference mel dump from disk so the encoder can be - // diffed against the reference WITHOUT introducing C++ mel - // drift into the comparison. Use this when a per-tensor - // compare regression fires and you want to isolate whether - // the drift originates in the C++ mel frontend or downstream - // in the encoder/decoder graph. validate.py exposes it via - // `--mel-from-ref`; the default validation path exercises the - // production C++ MelFrontend below so frontend regressions - // (e.g. precision changes that pass tolerance under ref-mel - // injection but break WER on borderline utterances) cannot - // slip through. + // 1. TRANSCRIBE_WHISPER_MEL_FROM_REF= — debug knob that reads the + // reference mel dump from disk so encoder drift can be isolated from + // C++ mel drift. The default path exercises the production frontend. // - // 2. C++ MelFrontend (default). Dual behavior: - // - Short-form (n_samples <= fe_n_samples): pad PCM to - // fe_n_samples (480000) and compute → exactly - // fe_nb_max_frames (3000) mel frames. Bit-identical to - // pre-Stage-2 behavior; tolerance + smoke tests exercise - // this branch. - // - Long-form (> fe_n_samples): compute mel on the raw audio - // (HF matches this — only short-form pads PCM to 30 s). The - // seek loop slices 3000-frame windows from the transposed - // buffer and zero-pads the trailing short window. + // 2. C++ MelFrontend (default): + // - Short-form (n_samples <= fe_n_samples): pad PCM to fe_n_samples + // (480000) -> exactly fe_nb_max_frames (3000) mel frames. + // - Long-form (> fe_n_samples): compute mel on raw audio (HF does the + // same — only short-form pads to 30s). The seek loop slices + // 3000-frame windows and zero-pads the trailing short window. const int n_mels = cm->hparams.enc_num_mel_bins; const int n_mel_frames_per_chunk = cm->hparams.fe_nb_max_frames > 0 ? cm->hparams.fe_nb_max_frames : 3000; const int n_samples_per_chunk = cm->hparams.fe_n_samples > 0 ? cm->hparams.fe_n_samples : 480000; - // Short-form = fits in a single 30 s window; bypass the dynamic - // stride + no-speech fallback machinery. This is a numeric AND - // behavioral parity gate: HF's generate() takes a different code - // path for is_shortform, and our Stage 1 tolerance tests were all - // captured in short-form. + // Short-form = fits a single 30s window; bypasses the dynamic-stride + // machinery. HF's generate() takes a different code path for is_shortform, + // so this is a numeric + behavioral parity gate (not just an optimization). const bool is_short_form = (n_samples <= n_samples_per_chunk); const int64_t t_mel_start = ggml_time_us(); @@ -1454,7 +1386,7 @@ transcribe_status whisper_run( return TRANSCRIBE_ERR_GGUF; } - // ----- Run-scoped constants -------------------------------------- + // ----- Run-scoped constants ----- const int64_t n_ctx_decoder = cm->hparams.dec_max_target_positions; const int64_t vocab_size = cm->hparams.dec_vocab_size; const int eos_id = cm->tok.eos_id() >= 0 ? cm->tok.eos_id() : 50257; @@ -1470,13 +1402,9 @@ transcribe_status whisper_run( requested_timestamps == TRANSCRIBE_TIMESTAMPS_AUTO || requested_timestamps == TRANSCRIBE_TIMESTAMPS_SEGMENT; - // Multilingual whisper variants emit <|lang|> + <|task|> tokens in - // the decoder prefix; English-only (.en) variants do not — their - // prefix is just <|sot|>. The two also differ at runtime: .en - // models do NOT carry <|translate|> / <|transcribe|> task tokens - // and have no per-language tokens to detect against. The - // supports_language_detect capability is the canonical signal; .en - // load sets it to false. + // Multilingual variants emit <|lang|> + <|task|> in the decoder prefix; + // .en variants have just <|sot|> and no translate/transcribe/language + // tokens. supports_language_detect is the canonical signal (.en = false). const bool is_multilingual = cm->caps.supports_language_detect; int32_t task_token = -1; @@ -1544,7 +1472,7 @@ transcribe_status whisper_run( return id >= timestamp_begin && id < static_cast(vocab_size); }; - // ----- Chunk loop ------------------------------------------------ + // ----- Chunk loop ----- cc->clear_result(); cc->chunk_traces.clear(); std::vector all_text_ids; @@ -1554,19 +1482,11 @@ transcribe_status whisper_run( int64_t t_decode_start = 0; bool t_decode_started = false; - // Finalize the context's result state from whatever accumulated so - // far. Called at normal completion AND at every mid-run abort exit - // (honoring the public contract in include/transcribe.h that - // partial segments/text accumulated before TRANSCRIBE_ERR_ABORTED - // must be readable via the normal accessors, distinguishable from - // "no result" by transcribe_was_aborted()). - // - // Matches end-of-run finalization: decodes all_text_ids to UTF-8, - // strips leading/trailing whitespace, pushes the text-only segment - // for TIMESTAMPS_NONE mode, and flips has_result. Safe to call - // from any point once the chunk loop has run at least zero times - // — an empty all_text_ids yields an empty text + an empty NONE-mode - // segment, which are the accessor sentinels anyway. + // Finalize the context's result state. Called at normal completion and at + // every mid-run abort exit (the public contract requires partial + // segments/text before TRANSCRIBE_ERR_ABORTED to stay readable via the + // normal accessors). Decodes all_text_ids to UTF-8, strips whitespace, + // pushes the text-only segment for TIMESTAMPS_NONE mode, flips has_result. auto commit_result = [&]() { cc->t_decode_us = t_decode_started ? ggml_time_us() - t_decode_start @@ -1613,38 +1533,18 @@ transcribe_status whisper_run( // Run-scoped Whisper run-ext pointer + RNG. // - // The whisper-specific knobs (initial prompt, prompt-token - // conditioning, temperature tuple, threshold checks, - // max_initial_timestamp cap) live on a kind-tagged family - // extension reached via transcribe_run_params::family. NULL selects - // transcribe_whisper_run_ext_init() — Whisper's own shipping recipe. - // The local `default_wp` must outlive the chunk loop because `wp` - // aliases it. - // - // The dispatcher has already validated that, if params->family is - // non-null, its header size+kind pass transcribe_model_accepts_ext_kind - // AND that the per-family minimum size passes whisper_run_validate - // (the pre-clear hook). We repeat transcribe_ext_check here as defense - // in depth right before casting the pointer; by this point it can only - // succeed, but keeping it makes the cast locally self-guarding. - // - // Pointer-field lifetimes (initial_prompt, prompt_tokens) are - // documented in include/transcribe/whisper.h: the library copies - // the referenced data into context state before transcribe_run - // returns; the caller may free the underlying buffers immediately - // after the call. Whisper realizes that contract by tokenizing - // initial_prompt into prev_ids below and by reading prompt_tokens - // into the same vector — neither pointer is retained. + // The whisper knobs live on a kind-tagged family extension reached via + // transcribe_run_params::family; NULL selects the shipping defaults from + // transcribe_whisper_run_ext_init(). `default_wp` must outlive the chunk + // loop because `wp` aliases it. transcribe_ext_check repeats the + // dispatcher's validation as defense in depth before casting. Pointer + // fields (initial_prompt, prompt_tokens) are copied into context state + // here, so the caller may free them right after the call. // - // RNG must be run-scoped, not chunk-scoped. HF does not reset its - // sampler between chunks (generation_whisper.py threads a single - // torch.Generator through the whole seek loop), so constructing - // std::mt19937 per chunk with the caller's seed would replay the - // same Mersenne-twister prefix at the start of each 30-second - // window — destroying determinism in exactly the case a fixed seed - // is meant to achieve (reproducible long-form transcripts). When - // seed == 0 we draw an OS entropy sample once and then let the - // single rng advance monotonically across chunks and tiers. + // RNG is run-scoped, not chunk-scoped: HF threads a single generator + // through the whole seek loop, so reseeding per chunk would replay the + // same prefix each window and destroy long-form determinism. seed == 0 + // draws OS entropy once, then the rng advances across chunks and tiers. if (const transcribe_status st = transcribe_ext_check( params != nullptr ? params->family : nullptr, TRANSCRIBE_EXT_KIND_WHISPER_RUN, @@ -1660,12 +1560,9 @@ transcribe_status whisper_run( : &default_wp; std::mt19937 rng(wp->seed != 0 ? wp->seed : std::random_device{}()); - // ===== Stage 3: prompt + condition_on_prev_tokens setup ============= - // - // HF generation_whisper.py:1743-1745 raises ValueError when - // prompt_condition_type=="all-segments" without - // condition_on_prev_tokens=True. Mirror as INVALID_ARG before any - // compute runs. + // Prompt + condition_on_prev_tokens setup. HF raises ValueError when + // prompt_condition_type=="all-segments" without condition_on_prev_tokens; + // mirror as INVALID_ARG before any compute runs. if (wp->prompt_condition == TRANSCRIBE_WHISPER_PROMPT_ALL_SEGMENTS && !wp->condition_on_prev_tokens) { @@ -1694,30 +1591,18 @@ transcribe_status whisper_run( return TRANSCRIBE_ERR_GGUF; } - // Resolve initial prompt → text-only token ids (no <|startofprev|>; - // the library prepends it). Two input paths: - // - // prompt_tokens: caller-owned bytes used verbatim. Caller must - // supply only the text-side tokens — we add <|startofprev|>. - // - // initial_prompt (string): tokenize HF's get_prompt_ids form - // ("<|startofprev|>" + " " + initial_prompt.strip()) and reject - // any special token (id >= eos_id) that appears in the text - // tokens, mirroring tokenization_whisper.py:716. We feed the - // leading-space-stripped form through the GPT-2 BPE encoder - // (which doesn't itself emit specials), so only special tokens - // present literally in user text would surface here. + // Resolve initial prompt -> text-only token ids (library prepends + // <|startofprev|>). Two paths: prompt_tokens (caller-owned, verbatim, + // text-side only); or initial_prompt string, tokenized as HF's + // get_prompt_ids form ("<|startofprev|> " + strip) with any special token + // (id >= eos_id) in the text rejected (tokenization_whisper.py). const int max_prev_cap = wp->max_prev_context_tokens > 0 ? wp->max_prev_context_tokens : (cm->hparams.dec_max_target_positions / 2 - 1); std::vector prompt_text_ids; if (wp->prompt_tokens != nullptr && wp->n_prompt_tokens > 0) { - // Caller-owned bytes used verbatim. The library prepends - // <|startofprev|>, so a leading prev_sot id from the caller - // would double-prepend — surface as INVALID_ARG so the - // mistake is caught immediately rather than silently producing - // a malformed prefix. (Header documents this contract; - // runtime check makes it observable.) + // The library prepends <|startofprev|>; a leading prev_sot id from the + // caller would double-prepend, so reject it as INVALID_ARG. if (prev_sot_id >= 0 && wp->prompt_tokens[0] == prev_sot_id) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper run: prompt_tokens must not include " @@ -1795,12 +1680,9 @@ transcribe_status whisper_run( prompt_text_ids.end() - max_prev_cap); } - // HF current_segments[0]. Store history as segment token slices, - // not one flattened vector, because _pad_to_max_length applies - // skip_ending_double_timestamps to each segment independently. - // For "first-segment" the prompt sits at the head of the history; - // for "all-segments" the history starts empty and the prompt is - // prepended per-chunk as the BOS. + // History stored as segment token slices (not one flat vector) because + // skip_ending_double_timestamps applies per-segment. FIRST_SEGMENT puts the + // prompt at the head; ALL_SEGMENTS starts empty and re-prepends per chunk. std::vector> prev_history_segments; if (wp->prompt_condition == TRANSCRIBE_WHISPER_PROMPT_FIRST_SEGMENT && !prompt_text_ids.empty()) @@ -1808,9 +1690,8 @@ transcribe_status whisper_run( prev_history_segments.push_back(prompt_text_ids); } - // Per-chunk decision; HF auto-disables on the previous chunk's hot - // accepted temperature (>= 0.5) per generation_whisper.py:1090-1093. - // Initialize to the user knob; flip per accepted tier at chunk end. + // Per-chunk; HF auto-disables when the previous chunk's accepted + // temperature was hot (>= 0.5). Init to the user knob; flipped at chunk end. bool do_condition_on_prev_tokens = wp->condition_on_prev_tokens; int seek = 0; @@ -1824,11 +1705,9 @@ transcribe_status whisper_run( const bool is_first_chunk = (seek == 0); const int64_t time_offset_ms = static_cast(seek) * 10; - // Slice + zero-pad the mel window to n_mel_frames_per_chunk. - // Zero-pad is ordered after the real frames, matching HF's - // F.pad(..., (0, num_segment_frames - cur_frames)). seek_num_frames - // is the count of real (unpadded) frames in this chunk — also the - // upper bound on how far `seek` can advance in one iteration. + // Slice + zero-pad the mel window to n_mel_frames_per_chunk (pad after + // the real frames, matching HF's F.pad). seek_num_frames is the real + // (unpadded) frame count and the max `seek` can advance per iteration. std::vector chunk_mel( static_cast(n_mels) * static_cast(n_mel_frames_per_chunk), @@ -1858,19 +1737,14 @@ transcribe_status whisper_run( t_decode_started = true; } - // Language detection — once, on first chunk, only if no hint - // and the model actually has language tokens. Non-multilingual - // (.en) models have no lang_token_ids so detection is skipped - // automatically; lang_token stays -1 and the prefix below omits - // the <|lang|> slot entirely. + // Language detection: once, first chunk, only if no hint and the model + // has language tokens (.en has none, so it's skipped automatically). if (is_first_chunk && is_multilingual && lang_token < 0 && !cm->lang_token_ids.empty()) { - // Lazily materialize encoder output to host for the - // prefill graph's encoder_out_in input. The hot path - // (cross-KV) reads cc->enc_out.tensor directly; this - // copy is one-shot per run when language detection - // is needed. + // Lazily materialize encoder output to host for the prefill + // graph's encoder_out_in input (one-shot per run; the hot path + // reads cc->enc_out.tensor directly). const int d_enc_lang = cc->enc_out.d_model; const int T_enc_lang = cc->enc_out.T_enc; cc->enc_host.resize(static_cast(d_enc_lang) * @@ -1925,10 +1799,8 @@ transcribe_status whisper_run( } } } - // Publish the detected ISO code only here, not in the - // user-hint branch above: this field's contract is "what the - // model told us when we asked it to pick," not "what the - // caller already knew." + // Publish the detected ISO code only here, not in the user-hint + // branch: the field means "what the model picked", not the hint. if (best_index >= 0 && static_cast(best_index) < cm->lang_codes.size()) { @@ -1942,28 +1814,16 @@ transcribe_status whisper_run( return TRANSCRIBE_ERR_GGUF; } - // ===== Stage 3: per-chunk prefix assembly ===================== - // - // Mirrors HF generation_whisper.py:_prepare_decoder_input_ids - // (1853-1918). The decoder input is prev_tokens + init_tokens: - // - // if do_condition_on_prev_tokens AND history non-empty: - // prev_tokens = bos + history[-cut_off:] - // bos = [<|startofprev|>, prompt_text_ids...] for ALL_SEGMENTS - // bos = [<|startofprev|>] for FIRST_SEGMENT - // skip_ending_double_timestamps: if history's last two ids are - // both timestamps, drop the last one (PR #34537 / #35750). - // elif initial prompt is set: - // prev_tokens = [<|startofprev|>, prompt_text_ids...] - // ALL_SEGMENTS: every chunk - // FIRST_SEGMENT: first chunk only - // else: - // prev_tokens = empty (Stage 2 behavior) - // - // HF (_prepare_decoder_input_ids) re-applies the prompt on - // EVERY chunk here regardless of prompt_condition. We deliberately - // diverge for FIRST_SEGMENT (the default) to match whisper.cpp / - // OpenAI, which prime only the first window. + // Per-chunk prefix assembly. Mirrors HF + // _prepare_decoder_input_ids: decoder input is prev_tokens + init. + // condition_on_prev + history: bos + history[-cut_off:], where bos is + // [<|startofprev|>(, prompt for ALL_SEGMENTS)]; drop a trailing + // double-timestamp per skip_ending_double_timestamps. + // else if initial prompt: [<|startofprev|>, prompt_text_ids...] + // (ALL_SEGMENTS every chunk; FIRST_SEGMENT first chunk only). + // else: empty. + // We diverge from HF for FIRST_SEGMENT (the default): prime only the + // first window, matching whisper.cpp / OpenAI. std::vector prev_tokens; if (do_condition_on_prev_tokens && !prev_history_segments.empty() && prev_sot_id >= 0) @@ -1980,10 +1840,8 @@ transcribe_status whisper_run( for (const auto & seg : prev_history_segments) { size_t n = seg.size(); if (n > 2 && token_is_timestamp(seg[n - 2])) { - // HF _pad_to_max_length(..., - // skip_ending_double_timestamps=True) drops the - // last token from any segment whose penultimate - // token is a timestamp. + // skip_ending_double_timestamps: drop the last token of any + // segment whose penultimate token is a timestamp. --n; } hist.insert(hist.end(), seg.begin(), seg.begin() + n); @@ -2005,9 +1863,8 @@ transcribe_status whisper_run( // Prefix for this chunk: // multilingual: prev_tokens + [SOT, lang, task, notimestamps?] // .en: prev_tokens + [SOT, notimestamps?] - // Non-multilingual models have neither <|lang|> nor <|task|> - // tokens in their vocab; emitting either would land on a - // garbage id. + // .en vocab has no <|lang|>/<|task|> tokens; emitting them would land + // on a garbage id. std::vector prompt_ids; prompt_ids.reserve(prev_tokens.size() + 4); prompt_ids.insert(prompt_ids.end(), @@ -2027,11 +1884,8 @@ transcribe_status whisper_run( // begin_index - start_of_trans_offset == len(prev_tokens)). const int sot_index = static_cast(prev_tokens.size()); - // Guard: prefix + 1 generated token must fit in the decoder - // self-attention window. With max_target_positions=448 and a - // maxed prev (224) + init (4) prefix this is 228 — plenty. A - // pathological caller passing prompt_tokens > 444 would land - // here. + // Guard: prefix + 1 generated token must fit the decoder self-attention + // window (only a pathological prompt_tokens > 444 trips this). if (seq_len + 1 > static_cast(n_ctx_decoder)) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "whisper run: prefix length %d exceeds decoder " @@ -2053,17 +1907,10 @@ transcribe_status whisper_run( } }; - // max_initial_timestamp_index: HF WhisperTimeStampLogitsProcessor - // (generation/logits_process.py:2036-2042) masks timestamps above - // `timestamp_begin + max_initial_timestamp_index` on the first - // generated token only. The default max_initial_timestamp is - // 1.0 s, which at Whisper's fixed 20 ms per timestamp token - // resolves to index 50 — i.e. the first emitted timestamp can - // mark a cut at most 1.0 s into the chunk. Prevents the decoder - // from skipping over an entire window's worth of audio on the - // opening cut. - // - // Negative / non-finite max_initial_timestamp disables the cap. + // max_initial_timestamp_index: HF WhisperTimeStampLogitsProcessor masks + // timestamps above timestamp_begin + this index on the first generated + // token only. Default 1.0s / 20ms-per-token = index 50, so the opening + // cut lands at most 1.0s in. Negative/non-finite disables the cap. const int max_initial_timestamp_index = (std::isfinite(wp->max_initial_timestamp) && wp->max_initial_timestamp >= 0.0f) @@ -2093,18 +1940,10 @@ transcribe_status whisper_run( int next_id = 0; - // ===== Stage 2.4: temperature fallback setup =================== - // - // Per-chunk temperature tuple = [t0, t0+dt, t0+2*dt, ...] up to - // and including 1.0. Default [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] — - // OpenAI whisper's shipping recipe. A tier is accepted when its - // compression_ratio < compression_ratio_thold AND its - // avg_logprob > logprob_thold. If all tiers fail, keep the last - // tier's output (HF does the same to avoid infinite loops). - // - // Thresholds use INF sentinels to mean "disabled"; see - // transcribe_whisper_run_ext_init() and the header's DISABLED - // constants. `wp` and `rng` are hoisted to run scope above. + // Temperature fallback setup. Per-chunk tuple [t0, t0+dt, ...] up to + // 1.0; default [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]. A tier is accepted when + // compression_ratio < thold AND avg_logprob > thold; if all fail, keep + // the last tier's output. Thresholds use INF sentinels to mean disabled. std::vector temperatures; temperatures.push_back(wp->temperature); if (wp->temperature_inc > 0.0f) { @@ -2125,38 +1964,23 @@ transcribe_status whisper_run( float accepted_avg_logprob = 0.0f; int accepted_n_fallbacks = 0; - // Stage 2.5 no-speech: probability captured from the SOT-position - // row of the prompt-pass logits (tier 0 only, matching HF's - // WhisperNoSpeechDetection which reruns on the full decoder - // prefix and reads logits[:, begin_index - start_of_trans_offset] - // — row 0 for our no-prev-context Stage 2 prefix). The token id - // is positional: no_speech = notimestamps - 1 (multilingual: - // 50362; .en: 50361). - // - // Post-decode, HF's _need_fallback - // (generation_whisper.py:1275-1286) sets should_skip=True when - // avg_logprob < logprob_thold AND - // no_speech_prob > no_speech_thold - // — both conditions required, not one alone. When should_skip - // fires, the chunk's output is discarded and seek advances by a - // full window (no fallback is attempted, which matches setting - // needs_fallback=False in HF). + // No-speech: probability captured from the SOT-position row of the + // prompt-pass logits (tier 0 only), matching HF WhisperNoSpeechDetection + // which reads logits[:, begin_index - start_of_trans_offset]. The token + // id is positional: no_speech = notimestamps - 1 (multilingual 50362, + // .en 50361). When the no-speech skip fires, the chunk's output is + // discarded and seek advances a full window (no fallback attempted). const int no_speech_token_id = cm->hparams.no_timestamps_token_id - 1; float no_speech_prob = 0.0f; bool no_speech_prob_captured = false; bool no_speech_fired_this_chunk = false; - // Port of HF _retrieve_compression_ratio's input convention - // (generation_whisper.py:1074-1086). The ratio is computed over - // the full generated tail including timestamp tokens, <|...|> - // specials, and EOS when present — fallback decides before EOS - // is stripped from seek_sequence. We mirror that here by - // assembling the metric input as: generated_ids + (EOS if the - // step loop terminated because it sampled eos_id). If the loop - // hit k_max_new_tokens before EOS, no EOS marker is appended — - // matching HF, whose seek_sequence in that case simply lacks an - // EOS. + // HF _retrieve_compression_ratio convention: the ratio is computed over + // the full generated tail including timestamps, specials, and EOS when + // present (fallback decides before EOS is stripped). We assemble the + // metric input as generated_ids + EOS-if-sampled; no EOS marker when + // the loop hit k_max_new_tokens first. bool tier_hit_eos = false; // KV cache allocation is tier-independent — done once per chunk @@ -2170,12 +1994,8 @@ transcribe_status whisper_run( cc->kv_cache.T_enc != T_enc_local; if (need_init) { cc->kv_cache.free(); - // AUTO: match cache dtype to the model's weight dtype. For - // F32 GGUFs that's F32 (zero-loss cache, no precision loss - // on the round-trip); for F16 and every quant preset it's - // F16 (the production default — weight+activation already - // pay F16 downcast costs, so the cache doesn't add new - // precision loss). Explicit --kv-type overrides. + // AUTO: match cache dtype to the weight dtype (F32 GGUF -> F32 + // zero-loss cache; F16 and quants -> F16). Explicit --kv-type wins. ggml_type kv_type_g; if (cc->kv_type == TRANSCRIBE_KV_TYPE_F32) { kv_type_g = GGML_TYPE_F32; @@ -2198,13 +2018,10 @@ transcribe_status whisper_run( } } - // ---- Cross-attention K/V precompute (chunk-scoped) ------------ - // - // Cross-K/V depends only on encoder output and the cross-attn - // weights — both are tier-invariant. Compute once per chunk and - // reuse across every fallback tier. Previously this lived inside - // the tier loop and re-ran every tier, paying n_layers × T_enc × - // d_model² of redundant matmul per tier on fallback chunks. + // ---- Cross-attention K/V precompute (chunk-scoped) ---- + // Cross-K/V depends only on the encoder output and cross-attn weights + // (both tier-invariant), so compute once per chunk and reuse across + // fallback tiers. { const int64_t t_cross_build_start = ggml_time_us(); if (!new_compute_ctx(8 * 1024 * 1024)) { @@ -2246,7 +2063,7 @@ transcribe_status whisper_run( cc->kv_cache.cross_populated = true; } - // ===== Tier loop (temperature fallback) ======================== + // ----- Tier loop (temperature fallback) ----- for (size_t ti = 0; ti < temperatures.size(); ++ti) { const float tier_T = temperatures[ti]; @@ -2261,9 +2078,8 @@ transcribe_status whisper_run( double sum_logprob = 0.0; int n_logprob_samples = 0; - // Dump gate. Tolerance references were captured at T=0 - // (tier 0 at defaults); only tier 0 of the first chunk - // emits intermediates. + // Dump gate: references were captured at T=0, so only tier 0 of the + // first chunk emits intermediates. const bool emit_tier_dumps = is_first_chunk && (ti == 0); auto tier_try_dump = [&](const char * name, ggml_tensor * t, const char * stage) @@ -2374,25 +2190,13 @@ transcribe_status whisper_run( const size_t row_bytes = static_cast(vocab_size) * sizeof(float); - // Capture no_speech_prob from the RAW SOT-position row, - // matching HF WhisperNoSpeechDetection semantics - // (transformers generation/logits_process.py:2095-2115): - // on the first generated-token step, when - // start_of_trans_offset > 1 — which is always our case - // (prefix always carries SOT + lang + task [+ - // notimestamps]) — HF reruns the model on the full - // decoder prefix and indexes - // logits[:, begin_index - start_of_trans_offset] - // which collapses to logits[:, sot_index] where sot_index - // is the position of <|startoftranscript|> within the - // prefix. With Stage 3 prev tokens this is len(prev_tokens); - // with no prev tokens it's 0 (matches the Stage 2 capture). - // - // Tier 0 only — matches HF's one-shot capture. Read the - // raw SOT row into a temporary buffer so suppress / - // begin_suppress / timestamp rules applied below to - // last_logits (the seq_len-1 row used to sample the first - // generated token) do not contaminate the measurement. + // Capture no_speech_prob from the RAW SOT-position row, matching + // HF WhisperNoSpeechDetection: logits[:, begin_index - + // start_of_trans_offset] collapses to logits[:, sot_index] + // (sot_index = position of <|startoftranscript|> in the prefix = + // len(prev_tokens), 0 with no prev tokens). Tier 0 only. Read + // into a temp buffer so the suppress / timestamp rules applied + // to last_logits below don't contaminate the measurement. if (ti == 0 && !no_speech_prob_captured && no_speech_token_id >= 0 && no_speech_token_id < static_cast(vocab_size)) @@ -2457,9 +2261,7 @@ transcribe_status whisper_run( ggml_time_us() - t_prompt_ts_start); if (tier_T <= 0.0f) { - // T==0: fused argmax + log-softmax. The whole call - // is recorded under sample; logprob bucket stays 0 - // for this path (it is computed for free here). + // T==0: fused argmax + log-softmax (logprob computed free). const int64_t t_prompt_sample_start = ggml_time_us(); float lp = 0.0f; next_id = sample_argmax_and_logprob(last_logits, &lp); @@ -2483,16 +2285,12 @@ transcribe_status whisper_run( cc->perf.prompt_cpu.add(ggml_time_us() - t_prompt_cpu_start); } - // ---- Step loop (n_tokens=1) ------------------------------- - // + // ---- Step loop (n_tokens=1) ---- // Two variants: - // GPU (Vulkan/Metal/CUDA/SYCL): build_step_graph — one - // static-topology graph per tier. KV writes via - // ggml_set_rows at runtime kv_idx; flash-attn reads a - // fixed max_n_kv window with a runtime mask. Removes - // per-step graph_build + sched_alloc. - // CPU (or debug, which needs the gen20 dump path): per- - // step build_decoder_graph_kv with a dynamic n_kv. + // GPU: build_step_graph — one static-topology graph per tier (KV + // writes via ggml_set_rows at runtime kv_idx, FA reads a fixed + // max_n_kv window with a runtime mask). + // CPU/debug: per-step build_decoder_graph_kv with dynamic n_kv. int n_past = seq_len; const bool primary_is_gpu = cm->plan.primary_kind != transcribe::BackendKind::Cpu && @@ -2501,18 +2299,15 @@ transcribe_status whisper_run( const bool use_step_graph = primary_is_gpu && !transcribe::debug::enabled(); - // Sized to fit prompt + max generated tail, padded to next - // pow2. For whisper this is typically 256 or 512 (cap is - // n_ctx_decoder=448). + // Sized to fit prompt + max generated tail, padded to next pow2, + // capped at n_ctx_decoder (448). int max_n_kv = 256; while (max_n_kv < seq_len + k_max_new_tokens) max_n_kv *= 2; if (max_n_kv > static_cast(n_ctx_decoder)) { max_n_kv = static_cast(n_ctx_decoder); } - // Step graph + persistent host buffers, only valid on - // use_step_graph path. Built lazily inside the loop on - // first use to keep the dynamic-graph fallback unaffected. + // Step graph + persistent host buffers (use_step_graph path only). StepBuild sb {}; std::vector step_mask; std::vector step_cross_mask; @@ -2638,10 +2433,8 @@ transcribe_status whisper_run( ggml_tensor * pos_in = ggml_graph_get_tensor(step_db.graph, "dec.pos_ids"); ggml_backend_tensor_set(pos_in, &pos, 0, sizeof(int32_t)); - // Padded step: build the trailing-padding mask. The - // valid range [0, n_past+1) is 0; trailing slots are - // -inf so the FA op ignores stale K/V values left in - // the cache from previous tiers / chunks. + // Padded step: [0, n_past+1) is 0, trailing slots -inf so + // FA ignores stale K/V from previous tiers/chunks. if (step_db.causal_mask_in != nullptr) { const int n_kv_mask = static_cast(step_db.causal_mask_in->ne[0]); @@ -2732,19 +2525,12 @@ transcribe_status whisper_run( cc->perf.step_cpu.add(ggml_time_us() - t_step_cpu_start); } - // ---- Compute tier metrics --------------------------------- - // - // Compression ratio runs over the full generated tail — all - // generated ids including timestamps / specials / EOS, with - // decoder_input_ids (the prompt header) stripped. This - // matches HF's _retrieve_compression_ratio input - // (generation_whisper.py:1074-1086 calls - // _retrieve_compression_ratio(seek_sequence, vocab_size) - // before EOS is stripped at :1085). An earlier iteration - // passed generated_text_ids here, which filtered out - // timestamps and specials and diverged from HF's metric by - // several percent on typical 20-50-token chunks — enough to - // flip accept/escalate on the fallback boundary. + // ---- Compute tier metrics ---- + // Compression ratio runs over the full generated tail (all ids + // including timestamps/specials/EOS, prompt header stripped), + // matching HF's _retrieve_compression_ratio. Using filtered + // text-only ids here would diverge enough to flip the fallback + // accept/escalate boundary. std::vector comp_tail; comp_tail.reserve(generated_ids.size() + 1); comp_tail.insert(comp_tail.end(), @@ -2759,46 +2545,27 @@ transcribe_status whisper_run( ? static_cast(sum_logprob / n_logprob_samples) : -std::numeric_limits::infinity(); - // Thresholds. - // - // HF _need_fallback (generation_whisper.py:1243-1287) falls - // back strictly on `comp_ratio > thold` or `avg_logprob < - // thold`; equality does NOT trigger fallback. So a tier is - // accepted (lp_ok/comp_ok) under `<= / >=`. The INF - // sentinels still work: tier_comp_ratio <= +INF is always - // true, tier_avg_logprob >= -INF is always true — disabled - // thresholds trivially satisfy the check. + // Thresholds. HF _need_fallback falls back strictly on + // `comp_ratio > thold` or `avg_logprob < thold`; equality does NOT + // trigger, so a tier is accepted under `<= / >=`. INF sentinels work + // (<= +INF and >= -INF are always true: disabled thresholds pass). const bool comp_ok = tier_comp_ratio <= wp->compression_ratio_thold; const bool lp_ok = tier_avg_logprob >= wp->logprob_thold; - // HF _need_fallback (generation_whisper.py:1275-1286) sets - // should_skip = True, needs_fallback = False - // when BOTH - // avg_logprob < logprob_threshold AND - // no_speech_prob > no_speech_threshold - // fire at the current tier. The skip halts fallback — HF - // explicitly flips needs_fallback to False so the outer - // generate_with_fallback loop terminates without trying - // hotter temperatures. no_speech_prob is a tier-0 capture - // (see above) so this is checked against the tier-0 value - // for every tier, but the logprob is per-tier. - // - // Disable sentinels: no_speech_thold = +INF makes the first - // comparison always false; logprob_thold = -INF makes the - // second always false. Either one disabled → no skip. + // HF _need_fallback sets should_skip (and halts fallback) when BOTH + // no_speech_prob > no_speech_thold AND avg_logprob < logprob_thold. + // no_speech_prob is the tier-0 capture (checked for every tier); + // logprob is per-tier. Sentinels: no_speech_thold=+INF or + // logprob_thold=-INF disables the skip. const bool no_speech_should_skip = no_speech_prob > wp->no_speech_thold && tier_avg_logprob < wp->logprob_thold; - // Commit the tier's output unconditionally. If no tier - // passes comp_ok/lp_ok AND no_speech_should_skip never - // fires, the last-tried tier wins (matches HF - // generate_with_fallback's behavior of keeping the final - // hypothesis when all tiers missed the thresholds). If - // no_speech_should_skip fires, we discard the tier output - // after recording the trace-visible metrics, below. + // Commit the tier's output unconditionally so the last-tried tier + // wins when no tier passes (matches HF generate_with_fallback). A + // no_speech_should_skip is discarded below after recording metrics. accepted_generated_ids = generated_ids; accepted_generated_text_ids = generated_text_ids; accepted_T = tier_T; @@ -2815,10 +2582,8 @@ transcribe_status whisper_run( } } - // Tier loop accepted something — hand off to segment emission. - // When no-speech fired, the tier produced output but we throw - // it away: seek still advances by a full window (per the - // _retrieve_segment branch in the seek-advance block below). + // Hand off the accepted tier to segment emission. A no-speech-fired + // chunk discards its output but still advances seek a full window. generated_ids = accepted_generated_ids; generated_text_ids = accepted_generated_text_ids; if (no_speech_fired_this_chunk) { @@ -2826,11 +2591,8 @@ transcribe_status whisper_run( generated_text_ids.clear(); } - // Stage 2.6 — commit per-chunk trace. Global window is - // [time_offset_ms, time_offset_ms + seek_num_frames*10ms). The - // trace captures what the fallback loop decided and what the - // no-speech gate saw; callers tracing a hallucination can map - // an output segment back to the tier/metric that produced it. + // Commit per-chunk trace over [time_offset_ms, +seek_num_frames*10ms): + // what the fallback loop decided and what the no-speech gate saw. { transcribe_whisper_chunk_trace trace; transcribe_whisper_chunk_trace_init(&trace); trace.t0_ms = time_offset_ms; @@ -2845,42 +2607,23 @@ transcribe_status whisper_run( cc->chunk_traces.push_back(trace); } - // Stage 3 — condition_on_prev_tokens gate for the NEXT chunk. - // - // HF generation_whisper.py:1090-1093 sets: - // is_low_temperature = temperature is None or temperature < 0.5 - // do_condition[i] = condition_on_prev_tokens AND is_low_temperature - // Hot tiers (>= 0.5) are presumed to have hallucinated, so the - // next chunk forgoes the prev-context window to avoid immediate - // propagation. The history itself is updated after - // _retrieve_segment below, because HF carries only the segment - // tokens that survive that slicing step. + // condition_on_prev_tokens gate for the NEXT chunk: HF disables it when + // the accepted temperature was hot (>= 0.5, presumed hallucinated) to + // avoid propagating into the next chunk. History is updated after + // _retrieve_segment (HF carries only the surviving segment tokens). do_condition_on_prev_tokens = wp->condition_on_prev_tokens && (accepted_T < 0.5f); - // ----- _retrieve_segment --------------------------------------- - // - // Port of HF generation_whisper.py:_retrieve_segment - // (1976-2073). Splits generated_ids on consecutive-timestamp - // pairs (each pair `<|t_a|><|t_b|>` marks "text between - // timestamps spans local time a..b"). Emits one SegmentEntry - // per pair. Returns segment_offset_frames — how far to advance - // `seek`: - // - // - Chunk ended with a single ts (pair not yet closed): no - // speech after last ts → advance past entire chunk - // (seek_num_frames). Preserves the tail for the next - // chunk's decoder. - // - Chunk has at least one closed pair but not single-ended: - // discard unfinished tail, advance by the last *closed* ts - // position × input_stride(2) mel frames. - // - No pairs at all: emit one segment covering the full - // chunk; advance by seek_num_frames. - // - // For TIMESTAMPS_NONE mode the prompt includes <|notimestamps|> - // so generated_ids carries no ts tokens → "no pairs" branch, - // full-chunk advance, no per-chunk segment emission (we gate - // on want_segment_timestamps). + // _retrieve_segment (port of HF). Splits generated_ids on consecutive- + // timestamp pairs (one SegmentEntry per pair) and returns + // segment_offset_frames (seek advance): + // - single trailing ts (pair not closed): advance the full chunk, + // preserving the tail for the next decoder. + // - >=1 closed pair, not single-ended: discard the unfinished tail, + // advance by the last closed ts position * input_stride(2) frames. + // - no pairs: emit one full-chunk segment, advance seek_num_frames. + // TIMESTAMPS_NONE includes <|notimestamps|>, so no ts tokens -> "no + // pairs" branch, full-chunk advance, no per-chunk segment emission. WhisperSegmentResult seg_res = whisper_retrieve_segment( generated_ids, cm->tok, time_offset_ms, seek_num_frames, want_segment_timestamps, timestamp_begin, vocab_size); @@ -2891,15 +2634,9 @@ transcribe_status whisper_run( std::vector> & prev_chunk_segments = seg_res.prev_chunk_segments; - // Stage 3 — update prev-context history for later chunks. - // - // HF appends the `segments` returned by _retrieve_segment to - // current_segments, then _prepare_decoder_input_ids builds the - // prev window from those segment tokens. In the closed-pair - // branch, _retrieve_segment intentionally discards unfinished - // tail tokens after the last completed timestamp pair; carrying - // accepted_generated_ids directly would leak that discarded - // tail into the next prompt and diverge from HF. + // Update prev-context history from the _retrieve_segment slices, not + // accepted_generated_ids directly: the closed-pair branch discards the + // unfinished tail, and carrying it would leak into the next prompt. if (!no_speech_fired_this_chunk && !prev_chunk_segments.empty()) { prev_history_segments.insert(prev_history_segments.end(), prev_chunk_segments.begin(), @@ -2911,25 +2648,13 @@ transcribe_status whisper_run( generated_text_ids.end()); // Seek advance. - // - // Short-form: fixed full-chunk advance. Because we padded PCM - // to exactly n_samples_per_chunk, the chunk's tail after the - // real audio is clamped-silence; the decoder would hallucinate - // "you you you" on it if we re-entered. HF's generate() avoids - // this by taking a separate is_shortform path that doesn't - // loop at all — we match by single-passing (a second iteration - // has nothing useful to decode). - // - // Long-form: dynamic advance per HF _retrieve_segment. The - // chunk's real content may not reach the 3000-frame window - // end; segment_offset_frames stops at the last closed - // timestamp pair so the next chunk picks up from there. - // Guarded against 0-advance infinite loops (fall back to full - // chunk). - // - // No-speech override (Stage 2.5): when both thresholds fire, - // the chunk is declared empty; skip the whole window so the - // next chunk decodes fresh encoder content. + // Short-form: full-chunk advance. PCM is padded to exactly + // n_samples_per_chunk, so the silent tail would hallucinate + // "you you you" if re-entered; single-pass it (matches HF's + // separate is_shortform path). + // Long-form: dynamic advance via segment_offset_frames (last closed + // timestamp pair), guarded against 0-advance infinite loops. + // No-speech fired: declared empty, skip the whole window. if (is_short_form) { seek += seek_num_frames; } else if (no_speech_fired_this_chunk) { @@ -2942,44 +2667,25 @@ transcribe_status whisper_run( } } - // Finalize: decode all_text_ids → UTF-8, strip whitespace, push the - // single text-only segment for TIMESTAMPS_NONE mode, flip - // has_result. Shared with the abort exit paths above so partial - // output is readable via the same accessors after - // TRANSCRIBE_ERR_ABORTED. commit_result(); return TRANSCRIBE_OK; } -// =========================================================================== -// Offline batched decode (transcribe_run_batch) -// =========================================================================== +// Offline batched decode (transcribe_run_batch). // -// Whisper is encoder-decoder with self + cross attention (the cohere / -// canary / moonshine pattern), so it batches the same way: the encoder -// stays SERIAL per utterance (compute-bound — the granite lesson) while the -// memory-bandwidth-bound autoregressive DECODE is batched across B -// utterances in lockstep, amortizing the per-step weight reads. +// Encoder stays SERIAL per utterance (compute-bound) while the bandwidth-bound +// autoregressive DECODE is batched across B utterances in lockstep. The batched +// fast path covers only the common, parity-exact case and peels everything else +// to the serial whisper_run(): +// batched: short-form (<= 30s) at tier 0 (T=0 greedy), TIMESTAMPS_NONE, no +// initial_prompt. +// serial : long-form, segment-timestamps, T > 0, prompt/prompt_tokens, +// invalid inputs, and any utterance whose tier-0 output fails the +// fallback thresholds (so the temperature ladder runs exactly). // -// Whisper's decode loop is far richer than its siblings (temperature -// fallback, timestamp rules, no-speech skip, long-form seek), so the -// batched fast path covers only the common, parity-exact case and PEELS -// everything else to the proven serial whisper_run(): -// -// batched : short-form (<= 30 s) utterances at tier 0 (temperature 0 = -// greedy argmax), TIMESTAMPS_NONE, no initial_prompt. -// serial : long-form, segment-timestamps, temperature > 0, prompt / -// prompt_tokens, invalid inputs, AND any batched utterance -// whose tier-0 output fails the fallback thresholds -// (compression-ratio / avg-logprob / no-speech-skip) — those -// re-run through whisper_run so the temperature ladder is -// applied exactly as in serial. -// -// For a clip that passes tier 0 (the overwhelming common case for clean -// speech) the batched greedy argmax is bit-for-bit the serial tier-0 -// result, so batch_parity holds; anything that would diverge is run -// serially and therefore trivially matches. +// A clip that passes tier 0 (the common case) is bit-for-bit the serial tier-0 +// result, so batch_parity holds; anything that would diverge runs serially. // Serial fallback: run each utterance through whisper_run and snapshot it. transcribe_status whisper_run_batch_serial( @@ -3664,28 +3370,17 @@ static bool whisper_accepts_ext_kind( return kind == TRANSCRIBE_EXT_KIND_WHISPER_RUN; } -// Pre-clear validation for the _RUN slot (see Arch::run_validate). Pure: -// it only inspects params->family. The dispatcher has already confirmed -// the ext header size and that WHISPER_RUN is accepted; here we enforce -// the per-kind minimum (the full transcribe_whisper_run_ext) before the -// previous result snapshot is cleared, so a too-small run ext is rejected -// without destroying the prior transcript. whisper_run repeats this check -// (defense in depth) right before it casts the pointer. -// -// Scope: this validates ext SHAPE only. Whisper's run-ext VALUE checks — -// prompt_condition=ALL_SEGMENTS without condition_on_prev_tokens, a -// prompt_tokens list that re-includes <|startofprev|>, and an -// initial_prompt carrying disallowed special tokens — live in whisper_run() -// and reject AFTER the snapshot is cleared, so a semantically-malformed (but -// correctly-sized) whisper run ext does clear the prior transcript. +// Pre-clear validation for the _RUN slot (see Arch::run_validate). Enforces the +// per-kind minimum (full transcribe_whisper_run_ext) before the prior result +// snapshot is cleared, so a too-small run ext is rejected without destroying +// the prior transcript. whisper_run repeats this (defense in depth) before the +// cast. // -// This is an intentional, accepted gap, not pending work. A malformed run -// ext is a caller programming error surfaced as INVALID_ARG; the run() path -// is one-shot (each call yields a fresh complete result the caller has -// already consumed), so unlike the streaming families there is no -// accumulating transcript to protect. The cost/benefit of hoisting the -// value checks here does not justify it. See docs/follow-ups.md for the -// full rationale should this ever need revisiting. +// Validates SHAPE only. The run-ext VALUE checks (ALL_SEGMENTS without +// condition_on_prev_tokens, prompt_tokens re-including <|startofprev|>, +// disallowed specials in initial_prompt) live in whisper_run() and reject after +// the snapshot is cleared — an accepted gap, since run() is one-shot with no +// accumulating transcript to protect. static transcribe_status whisper_run_validate( const transcribe_session * ctx, const transcribe_run_params * params) diff --git a/src/arch/whisper/public.cpp b/src/arch/whisper/public.cpp index 1fa282bc..36d71ce6 100644 --- a/src/arch/whisper/public.cpp +++ b/src/arch/whisper/public.cpp @@ -1,11 +1,9 @@ // arch/whisper/public.cpp - Whisper-family public C entry points. // -// The Whisper run-extension + chunk-trace init and accessor functions -// live here, not in the generic transcribe.cpp, so the central dispatcher -// stays family-agnostic (matching parakeet/moonshine, whose extension -// init functions live in their own family source). These symbols are part -// of the public ABI declared in include/transcribe/whisper.h and are -// linked into libtranscribe. +// The run-extension + chunk-trace init and accessor functions live here, not +// in the generic transcribe.cpp, so the central dispatcher stays family- +// agnostic (matching parakeet/moonshine). Public ABI from +// include/transcribe/whisper.h. #include "whisper.h" // WhisperSession (+ transcribe/whisper.h, transcribe-session.h) @@ -46,10 +44,7 @@ maybe_whisper_context(const struct transcribe_session * session) } // namespace -// --------------------------------------------------------------------------- // Init functions (single source of truth for Whisper struct defaults). -// --------------------------------------------------------------------------- - extern "C" void transcribe_whisper_chunk_trace_init(struct transcribe_whisper_chunk_trace * p) { if (p == nullptr) { return; } std::memset(p, 0, sizeof(*p)); @@ -73,10 +68,7 @@ extern "C" void transcribe_whisper_run_ext_init(struct transcribe_whisper_run_ex p->max_initial_timestamp = 1.0f; } -// --------------------------------------------------------------------------- // Chunk-trace accessors. -// --------------------------------------------------------------------------- - extern "C" int transcribe_get_whisper_chunk_count( const struct transcribe_session * session) { diff --git a/src/arch/whisper/weights.cpp b/src/arch/whisper/weights.cpp index c6eac588..36d8de50 100644 --- a/src/arch/whisper/weights.cpp +++ b/src/arch/whisper/weights.cpp @@ -1,29 +1,12 @@ -// arch/whisper/weights.cpp - implementation of read_whisper_hparams -// and build_whisper_weights. +// arch/whisper/weights.cpp - read_whisper_hparams + build_whisper_weights. +// Read every required KV explicitly, validate cross-field invariants, then +// resolve each tensor slot against the GGUF with expected shape. // -// Pattern mirrors cohere/weights.cpp: read every required KV -// explicitly, validate cross-field invariants, then resolve each -// tensor slot against the GGUF with expected shape. -// -// Whisper-specific notes: -// -// - Attention projections: q/v/out carry bias; k does NOT. The -// tensor catalog reflects this — no "attn.k.bias" slot exists -// (self-attn and cross-attn follow the same rule). -// -// - Logits head: no separate lm_head weight or bias. The converter -// does not emit dec.lm_head.weight; the decoder graph reuses -// dec.token_embd.weight via ggml_mul_mat at the final step. -// -// - suppress_tokens / begin_suppress_tokens are int32 arrays under -// the stt.whisper.* namespace. Empty vectors are the Absent case -// — the .en variants legitimately ship without these keys. -// -// - The converter tags normalization as "whisper_logmel", not one of -// the stat-based names (per_feature, global). This is the -// per-utterance log10(max(mel, 1e-10)) → clamp(max-8) → (+4)/4 -// dynamic-range compression. Any future frontend shared-helper -// must understand this tag. +// Whisper notes: q/v/out carry bias, k does NOT (no "attn.k.bias" slot, both +// self and cross); logits head has no separate lm_head, the decoder reuses +// dec.token_embd.weight (tied); suppress_tokens / begin_suppress_tokens are +// int32 arrays (empty for .en variants); normalize tag "whisper_logmel" = +// per-utterance log10(max(mel, 1e-10)) -> clamp(max-8) -> (+4)/4. #include "weights.h" @@ -106,10 +89,6 @@ transcribe_status read_optional_f32_kv(const gguf_context * gguf, } // namespace -// --------------------------------------------------------------------------- -// Hparams -// --------------------------------------------------------------------------- - transcribe_status read_whisper_hparams(const gguf_context * gguf, WhisperHParams & hp) { @@ -330,10 +309,7 @@ transcribe_status read_whisper_hparams(const gguf_context * gguf, return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // Frontend (mel filterbank + Hann window install) -// --------------------------------------------------------------------------- - transcribe_status install_mel_from_buffers( const WhisperHParams & hp, std::vector filterbank, @@ -374,10 +350,7 @@ transcribe_status install_mel_from_buffers( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // Weights -// --------------------------------------------------------------------------- - namespace { using transcribe::weights::lname; @@ -461,8 +434,7 @@ transcribe_status build_whisper_weights(ggml_context * ctx_meta, GET_LIN(b.attn_q_w, lname("enc.blocks.%d.attn.q.weight", i), d_model, d_model); GET_F32(b.attn_q_b, lname("enc.blocks.%d.attn.q.bias", i), d_model); - GET_LIN(b.attn_k_w, lname("enc.blocks.%d.attn.k.weight", i), d_model, d_model); - // attn.k has NO bias. + GET_LIN(b.attn_k_w, lname("enc.blocks.%d.attn.k.weight", i), d_model, d_model); // no bias GET_LIN(b.attn_v_w, lname("enc.blocks.%d.attn.v.weight", i), d_model, d_model); GET_F32(b.attn_v_b, lname("enc.blocks.%d.attn.v.bias", i), d_model); GET_LIN(b.attn_out_w, lname("enc.blocks.%d.attn.out.weight", i), d_model, d_model); @@ -478,9 +450,8 @@ transcribe_status build_whisper_weights(ggml_context * ctx_meta, // ----- decoder top ----- // Token embedding: PyTorch [vocab, dec_h] -> ggml ne = [dec_h, vocab]. - // This tensor doubles as the lm_head weight (tied), so it must - // accept the full quant allowlist (Embed bucket gets bumped to Q6_K - // under _M presets, F16 under F16, etc). + // Doubles as the tied lm_head weight, so it accepts the full quant + // allowlist (GET_LIN, not GET_F32). GET_LIN(weights.dec_top.token_embd_w, "dec.token_embd.weight", dec_h, vocab_size); GET_F32(weights.dec_top.pos_emb_w, "dec.pos_emb.weight", dec_h, tgt_pos); GET_F32(weights.dec_top.final_norm_w, "dec.final_norm.weight", dec_h); @@ -491,23 +462,23 @@ transcribe_status build_whisper_weights(ggml_context * ctx_meta, for (int i = 0; i < hp.dec_n_layers; ++i) { auto & b = weights.dec_blocks[i]; - // Self-attention (k has no bias). + // Self-attention. GET_F32(b.norm_self_w, lname("dec.blocks.%d.norm_self.weight", i), dec_h); GET_F32(b.norm_self_b, lname("dec.blocks.%d.norm_self.bias", i), dec_h); GET_LIN(b.self_q_w, lname("dec.blocks.%d.self_attn.q.weight", i), dec_h, dec_h); GET_F32(b.self_q_b, lname("dec.blocks.%d.self_attn.q.bias", i), dec_h); - GET_LIN(b.self_k_w, lname("dec.blocks.%d.self_attn.k.weight", i), dec_h, dec_h); + GET_LIN(b.self_k_w, lname("dec.blocks.%d.self_attn.k.weight", i), dec_h, dec_h); // no bias GET_LIN(b.self_v_w, lname("dec.blocks.%d.self_attn.v.weight", i), dec_h, dec_h); GET_F32(b.self_v_b, lname("dec.blocks.%d.self_attn.v.bias", i), dec_h); GET_LIN(b.self_out_w, lname("dec.blocks.%d.self_attn.out.weight",i), dec_h, dec_h); GET_F32(b.self_out_b, lname("dec.blocks.%d.self_attn.out.bias", i), dec_h); - // Cross-attention (k has no bias). + // Cross-attention. GET_F32(b.norm_cross_w, lname("dec.blocks.%d.norm_cross.weight", i), dec_h); GET_F32(b.norm_cross_b, lname("dec.blocks.%d.norm_cross.bias", i), dec_h); GET_LIN(b.cross_q_w, lname("dec.blocks.%d.cross_attn.q.weight", i), dec_h, dec_h); GET_F32(b.cross_q_b, lname("dec.blocks.%d.cross_attn.q.bias", i), dec_h); - GET_LIN(b.cross_k_w, lname("dec.blocks.%d.cross_attn.k.weight", i), dec_h, dec_h); + GET_LIN(b.cross_k_w, lname("dec.blocks.%d.cross_attn.k.weight", i), dec_h, dec_h); // no bias GET_LIN(b.cross_v_w, lname("dec.blocks.%d.cross_attn.v.weight", i), dec_h, dec_h); GET_F32(b.cross_v_b, lname("dec.blocks.%d.cross_attn.v.bias", i), dec_h); GET_LIN(b.cross_out_w, lname("dec.blocks.%d.cross_attn.out.weight",i), dec_h, dec_h); diff --git a/src/arch/whisper/weights.h b/src/arch/whisper/weights.h index 71c687df..679f1523 100644 --- a/src/arch/whisper/weights.h +++ b/src/arch/whisper/weights.h @@ -1,38 +1,10 @@ // arch/whisper/weights.h - canonical Whisper ASR tensor catalog and -// per-instance weight slots. +// per-instance weight slots. INTERNAL to src/arch/whisper/. // -// This header is INTERNAL to src/arch/whisper/. It defines: -// -// - WhisperHParams: the architecture KV the loader reads from -// stt.whisper.* / stt.frontend.* / stt.capability.* before -// allocating any tensors. Every dim that drives a tensor shape -// lives here. -// -// - WhisperWeights: a struct of named borrowed ggml_tensor* slots, -// one per logical weight in a Whisper ASR model. -// -// Architectural highlights: -// -// - Encoder: two Conv1d kernels (stride 1, then stride 2) forming -// the stem, then a learned positional embedding of fixed length -// max_source_positions (1500 for all whisper variants), then -// n_layers pre-LN transformer blocks with standard MHSA + FFN, -// then a final LayerNorm. -// -// - Decoder: learned token embedding + learned positional embedding -// (max_target_positions=448 for whisper-tiny), n_layers pre-LN -// blocks each with self-attention + cross-attention + FFN, final -// LayerNorm, logits head tied to the token embedding. -// -// - Whisper attention quirk: q_proj / v_proj / out_proj carry bias; -// k_proj does NOT (on both self- and cross-attention). The weight -// catalog encodes this by only declaring attn_k_w slots — no -// attn_k_b, no self_k_b, no cross_k_b. -// -// - Logits head: tied to dec.token_embd.weight; no separate bias. -// -// Tensor naming follows the Stage-3 convert-whisper.py contract -// (frontend.* / enc.* / dec.* as laid out in that script). +// WhisperHParams holds the architecture KV the loader reads (stt.whisper.* / +// stt.frontend.* / stt.capability.*); WhisperWeights holds the borrowed +// ggml_tensor* slots. Attention quirk: q/v/out carry bias, k does NOT (self +// and cross). Logits head is tied to dec.token_embd.weight, no separate bias. #pragma once @@ -50,10 +22,6 @@ struct ggml_tensor; namespace transcribe::whisper { -// --------------------------------------------------------------------------- -// Hyperparameters -// --------------------------------------------------------------------------- - struct WhisperHParams { // Encoder. int32_t enc_n_layers = 0; @@ -120,10 +88,6 @@ struct WhisperHParams { transcribe_status read_whisper_hparams(const gguf_context * gguf, WhisperHParams & hp); -// --------------------------------------------------------------------------- -// Weight slots -// --------------------------------------------------------------------------- - // Shared mel frontend buffers. preprocessor_config.json ships the // exact slaney filterbank and Hann window the model was trained with; // the converter stores them verbatim so C++ does not need to @@ -148,12 +112,10 @@ struct WhisperEncTop { ggml_tensor * final_norm_b = nullptr; // [d_model] }; -// One encoder transformer block. q/v/out have bias; k does NOT. +// One encoder transformer block. struct WhisperEncBlock { - // Pre-LN for self-attention. ggml_tensor * norm_attn_w = nullptr; ggml_tensor * norm_attn_b = nullptr; - // MHSA projections. ggml_tensor * attn_q_w = nullptr; ggml_tensor * attn_q_b = nullptr; ggml_tensor * attn_k_w = nullptr; // no bias @@ -161,10 +123,8 @@ struct WhisperEncBlock { ggml_tensor * attn_v_b = nullptr; ggml_tensor * attn_out_w = nullptr; ggml_tensor * attn_out_b = nullptr; - // Pre-LN for FFN. ggml_tensor * norm_ffn_w = nullptr; ggml_tensor * norm_ffn_b = nullptr; - // FFN (GELU). ggml_tensor * ffn_fc1_w = nullptr; ggml_tensor * ffn_fc1_b = nullptr; ggml_tensor * ffn_fc2_w = nullptr; @@ -180,9 +140,7 @@ struct WhisperDecTop { }; // One decoder block: self-attn + cross-attn + FFN, all pre-LN. -// q/v/out have bias; k does NOT (both self and cross). struct WhisperDecBlock { - // Pre-LN for self-attention. ggml_tensor * norm_self_w = nullptr; ggml_tensor * norm_self_b = nullptr; ggml_tensor * self_q_w = nullptr; @@ -193,7 +151,7 @@ struct WhisperDecBlock { ggml_tensor * self_out_w = nullptr; ggml_tensor * self_out_b = nullptr; - // Pre-LN for cross-attention (queries decoder state against encoder output). + // Cross-attention queries decoder state against encoder output. ggml_tensor * norm_cross_w = nullptr; ggml_tensor * norm_cross_b = nullptr; ggml_tensor * cross_q_w = nullptr; @@ -204,10 +162,8 @@ struct WhisperDecBlock { ggml_tensor * cross_out_w = nullptr; ggml_tensor * cross_out_b = nullptr; - // Pre-LN for FFN. ggml_tensor * norm_ffn_w = nullptr; ggml_tensor * norm_ffn_b = nullptr; - // FFN (GELU). ggml_tensor * ffn_fc1_w = nullptr; ggml_tensor * ffn_fc1_b = nullptr; ggml_tensor * ffn_fc2_w = nullptr; diff --git a/src/arch/whisper/whisper.h b/src/arch/whisper/whisper.h index 86ee6351..3686cce2 100644 --- a/src/arch/whisper/whisper.h +++ b/src/arch/whisper/whisper.h @@ -1,15 +1,6 @@ -// arch/whisper/whisper.h - Whisper ASR model and context types. -// -// This header is INTERNAL to src/arch/whisper/. It defines the concrete -// classes that derive from transcribe_model / transcribe_session for -// the Whisper ASR family (encoder-decoder transformer with cross- -// attention decoder). -// -// Shape is a direct descendant of src/arch/cohere/cohere.h: the two -// families share the encoder-decoder arch pattern, though their -// encoders are structurally different (cohere is Conformer, -// whisper is a stack of vanilla transformer blocks on top of a -// 2-layer conv stem). +// arch/whisper/whisper.h - Whisper ASR model and context types. INTERNAL to +// src/arch/whisper/. Defines the concrete transcribe_model / +// transcribe_session subclasses for the Whisper encoder-decoder family. #pragma once @@ -42,18 +33,10 @@ namespace transcribe::whisper { void apply_family_invariants(transcribe_model & model); -// --------------------------------------------------------------------------- -// KV cache for the autoregressive decoder. -// -// Self-attention KV cache: one entry per decode step, grows until EOS. -// Cross-attention KV cache: computed once from encoder output, reused -// across every decode step. -// -// Layout follows the cohere/whisper.cpp pattern: one flat 1D tensor -// per role; per-layer slices are constructed as views during graph -// build. -// --------------------------------------------------------------------------- - +// KV cache for the autoregressive decoder. Self-attention cache grows one +// entry per decode step until EOS; cross-attention cache is computed once from +// encoder output and reused across steps. One flat 1D tensor per role; +// per-layer slices are views built at graph time. struct WhisperKvCache { // Self-attention cache. // Flat tensors of size [d_model * n_layer * n_ctx]. @@ -125,21 +108,11 @@ struct WhisperKvCache { } }; -// --------------------------------------------------------------------------- -// Persistent encoder output tensor. -// -// Backend-resident F32 tensor sized [d_model, T_enc] that holds the -// encoder output across the encoder→cross-KV transition. Avoids the -// GPU→CPU→GPU roundtrip the previous design paid every chunk: the -// encoder graph appends a ggml_cpy into this tensor as its final op, -// and build_cross_kv_graph reads from it via a view. -// -// Lifetime: allocated on first use, sized from hparams.enc_d_model -// and the per-chunk T_enc the encoder emits. Reallocated only when -// shape changes (which does not happen for stock whisper variants — -// T_enc is fixed at hparams.enc_max_source_positions = 1500). -// --------------------------------------------------------------------------- - +// Persistent backend-resident F32 encoder output [d_model, T_enc], held across +// the encoder->cross-KV transition so cross-KV reads it via a view (no +// GPU->CPU->GPU roundtrip): the encoder graph ends with a ggml_cpy into it. +// Allocated on first use, reallocated only on shape change (T_enc is fixed at +// max_source_positions=1500 for stock variants, so effectively never). struct WhisperEncOut { ggml_tensor * tensor = nullptr; ggml_context * ctx = nullptr; @@ -167,17 +140,15 @@ bool enc_out_init(WhisperEncOut & enc_out, int d_model, int T_enc); -// Active-KV padding for the self-attention step graph. whisper.cpp's -// whisper_kv_cache_get_padding policy: 32 on Metal+FA, 1 otherwise -// (CUDA path uses 256 but is not exercised here yet). Padding the -// active n_kv lets the FA op land on a kernel that prefers shapes -// aligned to 32 — the trailing positions are masked to -inf so they -// do not contribute. Returns 1 (no padding) when use_flash is false. +// Active-KV padding for the self-attention step graph (whisper.cpp's +// whisper_kv_cache_get_padding): 32 on Metal+FA, 1 otherwise. Aligns the FA +// op's n_kv to a kernel-preferred multiple; trailing positions are masked to +// -inf. Returns 1 (no padding) when use_flash is false. int kv_pad_self_attn(transcribe::BackendKind kind, bool use_flash); -// Cross-attention cache padding multiple. Universal in whisper.cpp -// (not gated on backend); the cost is 36 unused rows per layer for -// whisper-tiny and the small cross-attn dilution that comes with it. +// Cross-attention cache padding multiple. Universal in whisper.cpp (not gated +// on backend); cost is a few unused rows per layer plus small cross-attn +// dilution. constexpr int k_cross_kv_pad = 256; // Allocate cache tensors. n_ctx caps self-attention length; T_enc is @@ -206,18 +177,8 @@ bool kv_cache_init_batched(WhisperKvCache & cache, int n_batch, ggml_type kv_type); -// --------------------------------------------------------------------------- -// Per-stage timing counters for performance diagnosis. -// -// Always-on (the timestamps themselves are negligible). Printing is -// opt-in via TRANSCRIBE_WHISPER_PROFILE=1 — see whisper_run / -// commit_result. Reset at the top of every whisper_run. -// -// Counters are organized by run-stage so a single summary line per -// stage shows the per-call breakdown of build / alloc / compute / -// tensor_get / cpu — matching the structural targets of Phases 2-4. -// --------------------------------------------------------------------------- - +// Per-stage timing counters. Always-on (timestamps are negligible); printing +// is opt-in via TRANSCRIBE_WHISPER_PROFILE=1. Reset at the top of whisper_run. struct WhisperPerfStage { int64_t total_us = 0; int count = 0; @@ -292,10 +253,6 @@ struct WhisperPerf { } }; -// --------------------------------------------------------------------------- -// Model / context -// --------------------------------------------------------------------------- - struct WhisperModel final : public transcribe_model { Tokenizer tok; WhisperHParams hparams; @@ -306,19 +263,15 @@ struct WhisperModel final : public transcribe_model { transcribe::BackendPlan plan; ggml_backend_buffer_t backend_buffer = nullptr; - // Language token ids keyed by BCP-47 short code read from - // general.languages. Populated at load time. Used by the decoder - // to resolve a params.language string to the correct - // <|lang_xx|> token id (whisper packs these in tokenizer slots - // 50259 + lang_index). + // Language token ids keyed by BCP-47 short code (from general.languages). + // Resolves a params.language string to the <|lang_xx|> token id (whisper + // packs these in tokenizer slots 50259 + lang_index). std::vector lang_codes; // owned copy; lifetime matches the model std::vector lang_token_ids; - // C++ mel frontend (per_utterance / hann_periodic / reflect / - // Slaney filterbank). Constructed at load() time using the - // filterbank + window baked into the GGUF for parity with the - // transformers reference. Optional so failures during load can - // still surface a usable model object for inspection. + // C++ mel frontend (per_utterance / hann_periodic / reflect / Slaney). + // Built from the filterbank + window baked into the GGUF. Optional so a + // load failure still surfaces a model object for inspection. std::optional mel; WhisperModel() = default; @@ -335,17 +288,12 @@ struct WhisperSession final : public transcribe_session { size_t compute_ctx_size = 0; ggml_backend_sched_t sched = nullptr; - // Persistent backend-resident encoder output. Encoder graph - // appends a ggml_cpy into enc_out.tensor as its final op so the - // cross-KV graph can read the encoder output via a view without - // a GPU→CPU→GPU roundtrip. + // Persistent backend-resident encoder output (see WhisperEncOut). WhisperEncOut enc_out; - // Host-side mirror of the encoder output, populated only when a - // path needs an F32 input fed via ggml_backend_tensor_set - // (currently: language detection on the first chunk and any - // dump-emitting validation runs). Kept off the per-chunk hot - // path so cross-KV no longer requires the materialization. + // Host-side mirror of the encoder output, populated only when a path needs + // an F32 input fed via ggml_backend_tensor_set (language detection on the + // first chunk, dump-emitting validation runs). Off the per-chunk hot path. std::vector enc_host; int enc_T = 0; // number of encoder frames (1500) @@ -355,31 +303,21 @@ struct WhisperSession final : public transcribe_session { // KV cache for the autoregressive decoder. WhisperKvCache kv_cache; - - // Flash-attention policy. Encoder runs at seqlen=1500 and always - // benefits from flash kernels where available; decoder runs at - // seqlen ≤ 448 but the cross-attention keys are 1500 long. Head - // dim is d_model/n_heads = 64 for tiny, which is supported by - // every backend's flash_attn_ext kernel today. Keep both on by - // default; TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH env vars - // still apply. + // Flash-attention policy. On by default for both; TRANSCRIBE_NO_FLASH / + // TRANSCRIBE_FORCE_FLASH still apply. bool encoder_use_flash = true; bool decoder_use_flash = true; - // Per-chunk decoding trace for the most recent successful run — - // populated by whisper_run as the chunk loop executes and exposed - // via transcribe_get_whisper_chunk_count / _get_whisper_chunk_trace. - // Cleared at the top of each run alongside cc->clear_result(). + // Per-chunk decoding trace for the most recent run, exposed via + // transcribe_get_whisper_chunk_count / _get_whisper_chunk_trace. Cleared + // at the top of each run alongside cc->clear_result(). std::vector chunk_traces; - // Per-stage timing counters. Always populated; a final summary is - // printed when TRANSCRIBE_WHISPER_PROFILE=1. + // Per-stage timing counters; summary printed when TRANSCRIBE_WHISPER_PROFILE=1. WhisperPerf perf; - // Reusable scratch for the multinomial T>0 sampler. Sized to - // vocab_size on first use; kept across steps to avoid the per-call - // double[vocab] allocation in the hot path. Argmax / T==0 do not - // touch this buffer. + // Reusable scratch for the multinomial T>0 sampler, sized to vocab_size on + // first use to avoid a per-call double[vocab] allocation in the hot path. std::vector sample_scratch; WhisperSession() = default; diff --git a/src/causal_lm/causal_lm.cpp b/src/causal_lm/causal_lm.cpp index 96a82275..b8464377 100644 --- a/src/causal_lm/causal_lm.cpp +++ b/src/causal_lm/causal_lm.cpp @@ -1,11 +1,5 @@ // src/causal_lm/causal_lm.cpp - shared causal-decoder LM block math -// (Llama / Qwen3 lineage; Qwen3 adds the per-head Q/K RMSNorm noted below). -// -// See causal_lm.h for the API contract. The block forward functions are -// strict extractions of the per-family code that lived in -// arch/qwen3_asr/decoder.cpp and arch/funasr_nano/decoder.cpp before -// this module landed; tensor topology, math, and KV memory layout are -// preserved bit-for-bit. +// (Llama / Qwen3 lineage). See causal_lm.h for the API contract. #include "causal_lm.h" #include "transcribe-session.h" @@ -24,21 +18,20 @@ namespace transcribe::causal_lm { namespace { -// Qwen3 RMSNorm: weight * rsqrt(mean(x²) + eps) * x. ggml_tensor * rms_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, float eps) { return ggml_mul(ctx, ggml_rms_norm(ctx, x, eps), weight); } -// Weight matmul forcing F32 accumulation for F16 weights. On CUDA, F16 weights -// take the cuBLAS path that accumulates in F16 (CUBLAS_COMPUTE_16F); a deep LLM -// residual stream has large-magnitude outlier channels that overflow F16's -// ~65504 range -> NaNs (multi-token prefill / batched GEMMs only; single-token -// decode uses an F32-accumulating mat-vec). GGML_PREC_F32 makes cuBLAS run the -// GEMM in F32, matching the CPU reference. Gated on F16 so BF16/quantized/F32 -// weights are untouched (BF16 is already COMPUTE_32F; quantized uses MMQ) and -// CPU/Metal — which always F32-accumulate — are unaffected. +// Weight matmul forcing F32 accumulation for F16 weights (CANONICAL F16-accum +// note for this module). On CUDA, F16 weights take the cuBLAS path that +// accumulates in F16 (CUBLAS_COMPUTE_16F); a deep LLM residual stream has +// large-magnitude outlier channels that overflow F16's ~65504 range -> NaNs. +// GGML_PREC_F32 makes cuBLAS run the GEMM in F32, matching the CPU reference. +// Gated on F16 so BF16/quantized/F32 weights are untouched (BF16 already +// COMPUTE_32F; quantized uses MMQ) and CPU/Metal (always F32-accumulate) are +// unaffected. ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * x) { ggml_tensor * y = ggml_mul_mat(ctx, w, x); @@ -50,10 +43,6 @@ ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * } // namespace -// --------------------------------------------------------------------------- -// KvCache -// --------------------------------------------------------------------------- - void KvCache::free() { if (buffer != nullptr) { ggml_backend_buffer_free(buffer); @@ -183,10 +172,7 @@ bool kv_init_batched(KvCache & cache, return true; } -// --------------------------------------------------------------------------- -// Block forward — prefill -// --------------------------------------------------------------------------- - +// Block forward — prefill. ggml_tensor * block_prefill( ggml_context * ctx, ggml_cgraph * gf, @@ -214,7 +200,6 @@ ggml_tensor * block_prefill( const size_t k_elem = ggml_element_size(kv_cache.self_k); const size_t v_elem = ggml_element_size(kv_cache.self_v); - // ---- Attention sub-layer ---- ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); // Q/K/V projections (bias-free on Qwen3). Packing into one mul_mat @@ -227,10 +212,9 @@ ggml_tensor * block_prefill( K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T_seq, 1); V = ggml_reshape_4d(ctx, V, head_dim, n_kv_heads, T_seq, 1); - // Per-head Q/K RMSNorm on head_dim (Qwen3 innovation). - // Per-head Q/K RMSNorm is a Qwen3 feature; Llama-style decoders - // (e.g. Voxtral's Ministral backbone) ship no q_norm/k_norm tensors, - // so the call site leaves these view slots null and we skip the norm. + // Per-head Q/K RMSNorm is a Qwen3 feature; Llama-style decoders ship no + // q_norm/k_norm tensors, so the call site leaves these slots null and we + // skip the norm. (Same for block_step / step_n / batched below.) if (view.attn_q_norm != nullptr) { Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, rms_eps), view.attn_q_norm); } @@ -253,13 +237,10 @@ ggml_tensor * block_prefill( rope_theta, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - // KV write: K, V have ne=[D, Hkv, T_seq, 1]; memory order - // (fastest→slowest) is D, Hkv, T. Cache stores position-major - // within each layer (per position, all D·Hkv values contiguous). - // Same layout — a 1D cpy handles it. - // Slab (layer, batch-slot) base offset, in elements. Defaults - // (kv_batch_slot=0, kv_n_batch=1) collapse to layer_idx*n_ctx*kv_dim — - // byte-identical to the single-shot layout. + // KV write: K, V have ne=[D, Hkv, T_seq, 1], same layout as the cache + // (position-major within each layer), so a 1D cpy handles it. Slab + // (layer, batch-slot) base offset in elements; defaults (kv_batch_slot=0, + // kv_n_batch=1) collapse to layer_idx*n_ctx*kv_dim. const size_t slab = static_cast(opts.kv_batch_slot) + static_cast(opts.kv_n_batch) * static_cast(layer_idx); @@ -338,8 +319,6 @@ ggml_tensor * block_prefill( o = mul_mat_f32acc(ctx, view.attn_o_w, o); x = ggml_add(ctx, x, o); - // Optional last-position slice before FFN (caller uses on the LAST - // block when only the final logits are consumed). if (opts.slice_last_before_ffn) { const int64_t hidden = x->ne[0]; const size_t elem = ggml_element_size(x); @@ -349,7 +328,6 @@ ggml_tensor * block_prefill( x = ggml_cont(ctx, x); } - // ---- MLP sub-layer (SwiGLU on packed gate_up) ---- ggml_tensor * ff_norm = rms_norm(ctx, x, view.norm_ffn_w, rms_eps); 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); @@ -360,10 +338,7 @@ ggml_tensor * block_prefill( return x; } -// --------------------------------------------------------------------------- -// Block forward — step -// --------------------------------------------------------------------------- - +// Block forward — step. ggml_tensor * block_step( ggml_context * ctx, ggml_cgraph * gf, @@ -392,7 +367,6 @@ ggml_tensor * block_step( const size_t k_elem = ggml_element_size(kv_cache.self_k); const size_t v_elem = ggml_element_size(kv_cache.self_v); - // ---- Attention sub-layer ---- ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); @@ -403,9 +377,6 @@ ggml_tensor * block_step( K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, 1, 1); V = ggml_reshape_4d(ctx, V, head_dim, n_kv_heads, 1, 1); - // Per-head Q/K RMSNorm is a Qwen3 feature; Llama-style decoders - // (e.g. Voxtral's Ministral backbone) ship no q_norm/k_norm tensors, - // so the call site leaves these view slots null and we skip the norm. if (view.attn_q_norm != nullptr) { Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, rms_eps), view.attn_q_norm); } @@ -451,9 +422,8 @@ ggml_tensor * block_step( } // Attention reads the FULL [0, max_n_kv) window. Mask zeros valid - // positions and -inf's the rest, so softmax ignores empty slots. - // Full-window read costs extra bandwidth (~1.6× avg) but the graph - // is static → zero per-step rebuild overhead. + // positions and -inf's the rest, so softmax ignores empty slots; the + // graph stays static (zero per-step rebuild). const size_t layer_off_bytes = k_elem * static_cast(layer_idx) * n_ctx * kv_dim; ggml_tensor * K_att = ggml_view_3d( @@ -505,7 +475,6 @@ ggml_tensor * block_step( o = mul_mat_f32acc(ctx, view.attn_o_w, o); x = ggml_add(ctx, x, o); - // ---- MLP sub-layer ---- ggml_tensor * ff_norm = rms_norm(ctx, x, view.norm_ffn_w, rms_eps); 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); @@ -516,10 +485,7 @@ ggml_tensor * block_step( return x; } -// --------------------------------------------------------------------------- -// Block forward — multi-position step (T_seq positions, single utterance) -// --------------------------------------------------------------------------- - +// Block forward — multi-position step (T_seq positions, single utterance). ggml_tensor * block_step_n( ggml_context * ctx, ggml_cgraph * gf, @@ -664,10 +630,7 @@ ggml_tensor * block_step_n( return x; } -// --------------------------------------------------------------------------- -// Block forward — batched step (B utterances) -// --------------------------------------------------------------------------- - +// Block forward — batched step (B utterances). ggml_tensor * block_step_batched( ggml_context * ctx, ggml_cgraph * gf, @@ -697,7 +660,6 @@ ggml_tensor * block_step_batched( const size_t k_elem = ggml_element_size(kv_cache.self_k); const size_t v_elem = ggml_element_size(kv_cache.self_v); - // ---- Attention sub-layer ---- ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); // [q_dim, B] @@ -710,9 +672,6 @@ ggml_tensor * block_step_batched( K = ggml_reshape_3d(ctx, K, head_dim, n_kv_heads, B); V = ggml_reshape_3d(ctx, V, head_dim, n_kv_heads, B); - // Per-head Q/K RMSNorm is a Qwen3 feature; Llama-style decoders - // (e.g. Voxtral's Ministral backbone) ship no q_norm/k_norm tensors, - // so the call site leaves these view slots null and we skip the norm. if (view.attn_q_norm != nullptr) { Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, rms_eps), view.attn_q_norm); } @@ -730,9 +689,8 @@ ggml_tensor * block_step_batched( GGML_ROPE_TYPE_NEOX, params.max_position, rope_theta, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - // ---- Batched KV write ---- - // Per layer the cache slab is [kv_dim, n_ctx, B]; one ggml_set_rows - // writes B rows (one per utterance) at B independent indices kv_idx[b]. + // Batched KV write: per layer the cache slab is [kv_dim, n_ctx, B]; one + // ggml_set_rows writes B rows at B independent indices kv_idx[b]. { const size_t layer_off_k = k_elem * static_cast(layer_idx) * n_ctx * kv_dim * B; @@ -755,7 +713,7 @@ ggml_tensor * block_step_batched( gf, ggml_set_rows(ctx, v_layer, V_row, kv_idx)); } - // ---- Attention read: [head_dim, max_n_kv, n_kv_heads, B] per layer ---- + // Attention read: [head_dim, max_n_kv, n_kv_heads, B] per layer. const size_t layer_off_bytes_k = k_elem * static_cast(layer_idx) * n_ctx * kv_dim * B; const size_t layer_off_bytes_v = @@ -792,7 +750,6 @@ ggml_tensor * block_step_batched( o = mul_mat_f32acc(ctx, view.attn_o_w, o); // [hidden, B] x = ggml_add(ctx, x, o); - // ---- MLP sub-layer ---- ggml_tensor * ff_norm = rms_norm(ctx, x, view.norm_ffn_w, rms_eps); 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); @@ -803,10 +760,7 @@ ggml_tensor * block_step_batched( return x; } -// --------------------------------------------------------------------------- -// Block forward — batched prefill (B utterances, T tokens each) -// --------------------------------------------------------------------------- - +// Block forward — batched prefill (B utterances, T tokens each). ggml_tensor * block_prefill_batched( ggml_context * ctx, ggml_cgraph * gf, @@ -837,7 +791,6 @@ ggml_tensor * block_prefill_batched( const size_t k_elem = ggml_element_size(kv_cache.self_k); const size_t v_elem = ggml_element_size(kv_cache.self_v); - // ---- Attention sub-layer ---- ggml_tensor * x_norm = rms_norm(ctx, x, view.norm_attn_w, rms_eps); ggml_tensor * Q = mul_mat_f32acc(ctx, view.attn_q_w, x_norm); // [q_dim, T, B] @@ -848,9 +801,6 @@ ggml_tensor * block_prefill_batched( K = ggml_reshape_4d(ctx, K, head_dim, n_kv_heads, T, B); V = ggml_reshape_4d(ctx, V, head_dim, n_kv_heads, T, B); - // Per-head Q/K RMSNorm is a Qwen3 feature; Llama-style decoders - // (e.g. Voxtral's Ministral backbone) ship no q_norm/k_norm tensors, - // so the call site leaves these view slots null and we skip the norm. if (view.attn_q_norm != nullptr) { Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, rms_eps), view.attn_q_norm); } @@ -866,7 +816,7 @@ ggml_tensor * block_prefill_batched( GGML_ROPE_TYPE_NEOX, params.max_position, rope_theta, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); - // ---- Batched KV write: each utterance's T rows -> [0, T) of its slab ---- + // Batched KV write: each utterance's T rows -> [0, T) of its slab. { const size_t layer_off_k = k_elem * static_cast(layer_idx) * n_ctx * kv_dim * B; @@ -889,7 +839,7 @@ ggml_tensor * block_prefill_batched( gf, ggml_set_rows(ctx, v_layer, V_rows, kv_idx)); } - // ---- Attention read: [head_dim, T, n_kv_heads, B] (positions [0, T)) ---- + // Attention read: [head_dim, T, n_kv_heads, B] (positions [0, T)). const size_t layer_off_bytes_k = k_elem * static_cast(layer_idx) * n_ctx * kv_dim * B; const size_t layer_off_bytes_v = @@ -920,7 +870,6 @@ ggml_tensor * block_prefill_batched( o = mul_mat_f32acc(ctx, view.attn_o_w, o); // [hidden, T, B] x = ggml_add(ctx, x, o); - // ---- MLP sub-layer ---- ggml_tensor * ff_norm = rms_norm(ctx, x, view.norm_ffn_w, rms_eps); 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); @@ -931,10 +880,6 @@ ggml_tensor * block_prefill_batched( return x; } -// --------------------------------------------------------------------------- -// pack_gate_up -// --------------------------------------------------------------------------- - void PackedGateUpHandles::free() { if (buffer != nullptr) { ggml_backend_buffer_free(buffer); @@ -977,9 +922,8 @@ bool pack_gate_up(ggml_backend_t backend, return false; } - // Allocate one packed tensor per block with the SAME dtype as - // gate_w. Mixed-quant blocks (e.g. all Q4_K_M) are fine because - // gate and up share a row-wise quant block size. + // One packed tensor per block, same dtype as gate_w (gate and up share + // a row-wise quant block size). for (size_t i = 0; i < entries.size(); ++i) { const auto & e = entries[i]; if (e.gate_w == nullptr || e.up_w == nullptr || @@ -1020,9 +964,8 @@ bool pack_gate_up(ggml_backend_t backend, ggml_backend_buffer_set_usage(out_handles.buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - // Fill: for row-wise quants (and f32/f16) concat-along-dim-1 is - // gate's bytes followed by up's bytes. One CPU round-trip per - // block at load time; zero cost at inference. + // Fill: concat-along-dim-1 is gate's bytes followed by up's bytes + // (one CPU round-trip per block at load time). std::vector buf; for (const auto & e : entries) { ggml_tensor * gate_up = *e.gate_up_w_out; @@ -1136,17 +1079,11 @@ transcribe_status run_batched_step_loop( stats->step_us = ggml_time_us() - t_step0; } - // A valid row was truncated if it stopped for a reason OTHER than emitting - // eos — i.e. it hit the generation budget (max_new) or the KV window. Note - // `finished` is NOT the discriminator: it is set on every stop reason (eos - // AND budget AND KV exhaustion, above), and the loop only exits once every - // valid row is finished, so `!finished[b]` is always false here. The actual - // signal is the last sampled token: `next_tok[b]` equals `eos_id` exactly - // when the row stopped by emitting eos (line `next_tok[b] = tok;` runs - // before the eos test, and `next_tok` is frozen once the row finishes), so - // `next_tok[b] != eos_id` means the row was cut off mid-transcript. This - // lets the family return per-utterance TRANSCRIBE_ERR_OUTPUT_TRUNCATED. See - // docs/input-limits.md. + // A valid row was truncated if it stopped for a reason OTHER than eos + // (generation budget or KV window). `finished` is set on every stop + // reason, so it can't discriminate; the signal is the last sampled token: + // `next_tok[b] != eos_id` means the row was cut off mid-transcript (it is + // frozen once the row finishes). See docs/input-limits.md. if (truncated_out != nullptr) { truncated_out->assign(n, 0); for (int b = 0; b < n; ++b) { diff --git a/src/causal_lm/causal_lm.h b/src/causal_lm/causal_lm.h index e17e2a35..a7935722 100644 --- a/src/causal_lm/causal_lm.h +++ b/src/causal_lm/causal_lm.h @@ -1,65 +1,16 @@ // src/causal_lm/causal_lm.h - shared causal-decoder LM block math, // KV cache, and packed gate/up MLP packing. // -// This is the dense decoder-only transformer block of the Llama / Qwen3 -// lineage: pre-LN RMSNorm (model-provided eps via BlockParams::rms_eps), -// GQA, NeoX rotate_half RoPE, SwiGLU MLP via packed gate+up, with -// OPTIONAL per-head Q/K RMSNorm. -// It is NOT Qwen3-specific despite the families that first used it — the -// optional slots and `BlockParams` cover several backbones: -// -// - Qwen3 decoder (per-head Q/K RMSNorm present, θ=1e6 NeoX RoPE; -// text-only MRoPE collapses to NeoX): used by arch/qwen3_asr -// (`Qwen3ASRThinkerTextModel`), arch/funasr_nano (bundled -// `Qwen3-0.6B`), and arch/canary_qwen (`Qwen3-1.7B`). These pass -// attn_q_norm / attn_k_norm. -// - Llama / Ministral decoder (no per-head Q/K norm): used by -// arch/voxtral and arch/voxtral_realtime, which leave the q/k-norm -// slots null. voxtral_realtime additionally drives a per-layer -// delay-conditioned FFN scale through the optional `ffn_scale` slot. -// - Infra-only reuse: arch/granite reuses KvCache, pack_gate_up, and -// run_batched_step_loop but keeps its OWN block math (Granite's -// attention/residual/embedding/logits multipliers don't fit the -// shared block). See arch/granite/decoder.cpp. -// -// All callers pack ffn_gate_w + ffn_up_w into a single -// [hidden, 2·intermediate] tensor at load time so the FFN can run with -// one mul_mat instead of two. -// -// This module exposes those primitives via the same view + -// free-function pattern used by `src/sanm/` and `src/conformer/`: -// -// - BlockView : nullable-pointer projection over one decoder -// block's weights. Each family fills it at the -// call site from its own `DecBlock` struct. -// - BlockParams : per-call scalar shape (n_heads, n_kv_heads, -// head_dim, rms_eps, rope_theta, max_position). -// - block_prefill / block_step : the two block-forward -// variants used by the prefill (T_seq > 1) and -// autoregressive step (T_seq == 1) graphs. -// - KvCache + kv_init : self-attention-only KV cache sized -// [n_kv_heads · head_dim · n_ctx · n_layer] flat -// per K and V. Layout matches what both families -// emitted before extraction (so the conversion is -// a strict refactor — no graph topology change). -// - pack_gate_up : load-time helper that allocates the -// [hidden, 2·intermediate] packed tensor per -// block and copies the gate and up bytes in -// (compatible with row-wise quants because -// concat-along-dim-1 is just bytes-concat for -// those types). -// -// What deliberately stays per-family (audio-LLM call sites): -// -// - Token embedding lookup and audio injection (the encoder/adaptor -// output shape and splice point differ per family), the final norm -// and lm_head (tied for the Qwen3 families, UNTIED for voxtral, and -// Granite additionally scales embeddings/logits), dump-tensor naming -// for validate.py parity, prefill / step graph allocation, the -// autoregressive driver loop, and the chat-template / special-token -// handling. The block-forward helpers below return raw `ggml_tensor*` -// so each family keeps full control of its dump points and tensor -// names. +// Dense decoder-only transformer block of the Llama / Qwen3 lineage: +// pre-LN RMSNorm, GQA, NeoX rotate_half RoPE, SwiGLU MLP via packed +// gate+up, with OPTIONAL per-head Q/K RMSNorm (Qwen3) and an optional +// per-layer FFN scale (voxtral_realtime). Not Qwen3-specific: the +// optional slots cover Qwen3, Llama/Ministral, and infra-only reuse +// by Granite (KvCache / pack_gate_up / run_batched_step_loop only). +// Primitives follow the view + free-function pattern of src/sanm/ and +// src/conformer/. Helpers return raw ggml_tensor* so each family owns +// embedding/audio-injection, final-norm/lm_head, dump naming, graph +// allocation, the decode driver, and chat-template handling. #pragma once @@ -75,19 +26,10 @@ struct transcribe_session; namespace transcribe::causal_lm { -// --------------------------------------------------------------------------- -// Block view -// --------------------------------------------------------------------------- - // Nullable-pointer projection over one causal-decoder block's weights. -// Field names match `arch/qwen3_asr/QwenAsrDecBlock` and -// `arch/funasr_nano/DecBlock` so the call-site builder is a struct -// initializer. -// -// All slots are required for block_prefill / block_step EXCEPT three -// optional ones (null = skip): attn_q_norm / attn_k_norm (per-head Q/K -// norm, absent on Llama-style decoders) and ffn_scale (per-layer FFN -// scale, used only by voxtral_realtime). See each field's note below. +// Field names match arch/qwen3_asr and arch/funasr_nano so the call-site +// builder is a struct initializer. All slots required except three +// optional ones (null = skip): attn_q_norm / attn_k_norm and ffn_scale. struct BlockView { ggml_tensor * norm_attn_w = nullptr; // input_layernorm (RMSNorm, no bias) ggml_tensor * norm_ffn_w = nullptr; // post_attention_layernorm @@ -95,21 +37,17 @@ struct BlockView { ggml_tensor * attn_k_w = nullptr; // [hidden, n_kv_heads * head_dim] ggml_tensor * attn_v_w = nullptr; // [hidden, n_kv_heads * head_dim] ggml_tensor * attn_o_w = nullptr; // [n_heads * head_dim, hidden] - // Per-head Q/K RMSNorm (Qwen3). OPTIONAL: leave null for Llama-style - // decoders (e.g. Voxtral's Ministral backbone) that ship no q/k norm; - // the block helpers skip the norm when the slot is null. + // Per-head Q/K RMSNorm (Qwen3). Null on Llama-style decoders (e.g. + // 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: [hidden, 2·intermediate]. Filled by pack_gate_up - // at load time (see below). The graph runs ONE mul_mat against this - // tensor + ggml_swiglu instead of two mul_mats + manual silu·mul. + // 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. ggml_tensor * ffn_gate_up_w = nullptr; // [hidden, 2·intermediate] ggml_tensor * ffn_down_w = nullptr; // [intermediate, hidden] - // OPTIONAL per-layer FFN-branch scale applied right after norm_ffn and - // before the gate_up mul_mat: ff_norm *= ffn_scale (broadcast [hidden,1] - // over the token axis). Null for every standard caller; used by - // voxtral_realtime's delay-conditioned adaptive-norm (a per-layer constant - // `1 + ada(t_cond)` precomputed once per run). See arch/voxtral_realtime. + // Optional per-layer FFN-branch scale (ff_norm *= ffn_scale, broadcast + // over the token axis). Null for standard callers; voxtral_realtime's + // delay-conditioned adaptive-norm uses it. ggml_tensor * ffn_scale = nullptr; // [hidden] or [hidden,1], or null }; @@ -123,13 +61,9 @@ struct BlockParams { float rope_theta = 0.0f; }; -// --------------------------------------------------------------------------- -// KV cache (self-attention only) -// --------------------------------------------------------------------------- - -// Flat 1D K and V tensors sized [n_kv_heads · head_dim · n_ctx · n_layer]. -// Layout (slowest → fastest): layer, position, head, dim. Matches the -// memory the per-family decoder.cpp wrote before extraction. +// KV cache (self-attention only). Flat 1D K and V tensors sized +// [n_kv_heads · head_dim · n_ctx · n_layer]. Layout (slowest → fastest): +// layer, position, head, dim. struct KvCache { ggml_tensor * self_k = nullptr; ggml_tensor * self_v = nullptr; @@ -140,9 +74,8 @@ struct KvCache { int n_ctx = 0; int n = 0; // current fill (high-water mark) int head = 0; // next write position - int n_batch = 1; // utterances packed along the batch axis (offline - // batched decode). 1 == single-shot layout, in which - // case the flat tensor is byte-identical to kv_init. + int n_batch = 1; // utterances packed along the batch axis + // (offline batched decode); 1 == single-shot void free(); }; @@ -159,8 +92,7 @@ bool kv_init(KvCache & cache, // Batched variant: allocates [n_kv_heads · head_dim · n_ctx · n_batch · n_layer] // flat per K and V, laid out (slowest → fastest) layer, batch, position, head, -// dim. With n_batch == 1 the size and layout are identical to kv_init, so the -// single-shot block_step views read byte-identical memory. Sets cache.n_batch. +// dim. With n_batch == 1 the size and layout match kv_init. Sets cache.n_batch. bool kv_init_batched(KvCache & cache, ggml_backend_t backend, int n_ctx, @@ -170,40 +102,30 @@ bool kv_init_batched(KvCache & cache, int n_batch, ggml_type kv_type); -// --------------------------------------------------------------------------- -// Block forward (prefill: T_seq > 1) -// --------------------------------------------------------------------------- - struct BlockOpts { bool use_flash = true; - // When true, after the post-attention residual add but before the - // FFN sub-layer, slice x to the last position only. Output shape - // becomes [hidden, 1] instead of [hidden, T_seq]. The caller uses - // this on the LAST block when only the final position's logits are - // consumed (matches llama.cpp's inp_out_ids optimization). + // When true, after the post-attention residual add but before the FFN, + // slice x to the last position only (output [hidden, 1]). Used on the + // LAST block when only the final logits are consumed (llama.cpp's + // inp_out_ids optimization). bool slice_last_before_ffn = false; // Offline batched decode: write/read this utterance's KV in slab // `kv_batch_slot` of a batched cache holding `kv_n_batch` slabs. The // per-(layer,slot) byte offset becomes // (kv_batch_slot + kv_n_batch * layer_idx) * n_ctx * kv_dim. - // Defaults (0, 1) reproduce the single-shot offset layer_idx*n_ctx*kv_dim - // exactly, so single-shot callers are unaffected. + // Defaults (0, 1) reproduce the single-shot offset. int kv_batch_slot = 0; int kv_n_batch = 1; }; -// Run one causal-LM block on `x` (ne = [hidden, T_seq]). Writes K/V for -// positions [0, T_seq) at layer `layer_idx` of `kv_cache`. Mask is -// [T_seq, T_seq] f16 (causal, host-prepared); positions are [T_seq] i32. -// -// Returns the post-block hidden state. Shape is [hidden, T_seq] unless -// `opts.slice_last_before_ffn`, in which case [hidden, 1]. -// -// The function does NOT name any tensors with ggml_set_name and does -// NOT call ggml_set_output / mark_tensor_for_dump. The family's prefill -// builder owns naming and dump-point selection. +// Block forward (prefill, T_seq > 1). Runs one block on `x` +// (ne = [hidden, T_seq]), writing K/V for positions [0, T_seq) at layer +// `layer_idx`. Mask is [T_seq, T_seq] f16 causal; positions are [T_seq] +// i32. Returns the post-block hidden state ([hidden, T_seq], or +// [hidden, 1] under opts.slice_last_before_ffn). Names no tensors; the +// family's prefill builder owns naming and dump-point selection. ggml_tensor * block_prefill( ggml_context * ctx, ggml_cgraph * gf, @@ -217,18 +139,12 @@ ggml_tensor * block_prefill( ggml_tensor * positions, // [T_seq] i32 BlockOpts opts = {}); -// --------------------------------------------------------------------------- -// Block forward (step: T_seq == 1) -// --------------------------------------------------------------------------- - -// Run one causal-LM block on a single new token. Writes K/V for the row -// indexed by `kv_idx` (i64 [1]) into layer `layer_idx`. Reads the full -// [0, max_n_kv) KV window for attention; `mask` is [max_n_kv, 1] f16 -// with zeros at positions ≤ n_past and -inf beyond (host-prepared). -// -// `position` is [1] i32 (RoPE), `kv_idx` is [1] i64 (set_rows index). -// Both are equal at runtime but kept distinct because RoPE wants i32 -// and set_rows wants i64. +// Block forward (step, T_seq == 1). Runs one block on a single new token, +// writing K/V for the row indexed by `kv_idx` (i64 [1]) into layer +// `layer_idx`. Reads the full [0, max_n_kv) KV window; `mask` is +// [max_n_kv, 1] f16 (zeros ≤ n_past, -inf beyond). `position` ([1] i32) +// and `kv_idx` ([1] i64) are equal at runtime but kept distinct (RoPE +// wants i32, set_rows wants i64). ggml_tensor * block_step( ggml_context * ctx, ggml_cgraph * gf, @@ -243,15 +159,11 @@ ggml_tensor * block_step( ggml_tensor * kv_idx, // [1] i64 bool use_flash); -// --------------------------------------------------------------------------- -// Block forward (multi-position step: T_seq new tokens in one pass) -// --------------------------------------------------------------------------- - -// Same as block_step but processes T_seq positions in a single forward. -// Writes T_seq rows at the indices in `kv_idx` (i64 [T_seq]) and reads the -// full [0, max_n_kv) KV window. `mask` is [max_n_kv, T_seq] f16 — column -// `c` masks the c-th query position (typically causal: 0 for cached -// positions <= cur_pos+c, -inf beyond). Used by spec-decode verify pass. +// Block forward (multi-position step). Like block_step but processes +// T_seq positions in one forward: writes T_seq rows at the indices in +// `kv_idx` (i64 [T_seq]) and reads the full [0, max_n_kv) window. `mask` +// is [max_n_kv, T_seq] f16 — column `c` masks the c-th query position. +// Used by the spec-decode verify pass. ggml_tensor * block_step_n( ggml_context * ctx, ggml_cgraph * gf, @@ -267,19 +179,13 @@ ggml_tensor * block_step_n( ggml_tensor * kv_idx, // [T_seq] i64 bool use_flash); -// --------------------------------------------------------------------------- -// Block forward (batched step: B utterances, one new token each) -// --------------------------------------------------------------------------- - -// Run one causal-LM block on B new tokens (x = [hidden, B]), one per utterance -// stepping in lockstep against a batched KV cache (kv_init_batched, n_batch=B). -// -// Layout choice: the batch axis sits on ne[2] (the "position/token" axis) -// through projection + RoPE so each utterance gets its OWN RoPE position from -// `position` [B]; Q is then permuted to [head_dim, 1, n_heads, B] for -// flash_attn_ext, which batches on ne[3]. K/V are written with a single -// batched ggml_set_rows (dest [kv_dim, n_ctx, B], src [kv_dim, 1, B], idx -// [1, B]) so utterance b writes row kv_idx[b] of its own slab. +// Block forward (batched step). Runs one block on B new tokens +// (x = [hidden, B]), one per utterance stepping in lockstep against a +// batched KV cache (kv_init_batched, n_batch=B). The batch axis rides +// ne[2] through projection + RoPE so each utterance gets its own RoPE +// position from `position` [B]; Q is then permuted to [head_dim, 1, +// n_heads, B] for flash_attn_ext (batches on ne[3]). One batched +// ggml_set_rows writes B KV rows. // // Inputs: // x [hidden, B] @@ -290,8 +196,7 @@ ggml_tensor * block_step_n( // kv_idx [1, B] i64 — KV write row per utterance (= that row's n_past) // // Requires use_flash (the manual GQA path is single-shot only). Returns -// [hidden, B]. With B == 1 this is numerically the batched-of-one analog of -// block_step; the parity gate compares batched row b vs single-shot b. +// [hidden, B]. ggml_tensor * block_step_batched( ggml_context * ctx, ggml_cgraph * gf, @@ -307,21 +212,15 @@ ggml_tensor * block_step_batched( ggml_tensor * kv_idx, // [1, B] i64 bool use_flash); -// --------------------------------------------------------------------------- -// Block forward (batched prefill: B utterances, T tokens each) -// --------------------------------------------------------------------------- - -// Run one causal-LM block over B prompts of T tokens each (x = [hidden, T, B]), -// writing each utterance's K/V to positions [0, T) of its own slab in a -// batched KV cache (kv_init_batched, n_batch=B). Batch rides ne[2]; RoPE -// positions [T] are shared (every prompt starts at position 0); the causal -// mask [T, T] is shared (a real token p < T_prompt[b] only attends [0, p], -// all real). One batched ggml_set_rows writes B*T KV rows (idx[t,b]=t). -// -// Pad tokens (t >= T_prompt[b]) compute harmlessly: their KV lands in pad -// slab positions that the step loop overwrites before attending, and the -// caller gathers only each utterance's real last-position output. Requires -// use_flash (the manual GQA path has no batch axis). Returns [hidden, T, B]. +// Block forward (batched prefill). Runs one block over B prompts of T +// tokens each (x = [hidden, T, B]), writing each utterance's K/V to +// positions [0, T) of its own slab (kv_init_batched, n_batch=B). Batch +// rides ne[2]; RoPE positions [T] and the causal mask [T, T] are shared. +// One batched ggml_set_rows writes B*T KV rows (idx[t,b]=t). Pad tokens +// (t >= T_prompt[b]) compute harmlessly: their KV lands in pad positions +// the step loop overwrites before attending, and the caller gathers only +// each utterance's real last-position output. Requires use_flash. +// Returns [hidden, T, B]. ggml_tensor * block_prefill_batched( ggml_context * ctx, ggml_cgraph * gf, @@ -337,15 +236,11 @@ ggml_tensor * block_prefill_batched( ggml_tensor * kv_idx, // [T_seq, B] i64 (idx[t,b] = t) bool use_flash); -// --------------------------------------------------------------------------- -// Load-time gate/up packing -// --------------------------------------------------------------------------- +// Load-time gate/up packing. // One block's input pointers (gate_w, up_w) and output slot -// (gate_up_w_out, written by pack_gate_up). gate_w must already be -// allocated and bound to a backend buffer; up_w likewise. The output -// pointer is written with the address of the newly-allocated packed -// tensor inside handles.ctx. +// (gate_up_w_out, written by pack_gate_up). gate_w / up_w must already be +// allocated and bound to a backend buffer. struct GateUpEntry { ggml_tensor * gate_w; ggml_tensor * up_w; @@ -364,20 +259,13 @@ struct PackedGateUpHandles { void free(); }; -// Allocate a separate packed context + a [hidden, 2·intermediate] -// tensor per block on `backend`, then copy gate's bytes followed by -// up's bytes into each packed tensor. Compatible with row-wise quants -// (Q4/Q5/Q6/Q8) because concat-along-dim-1 is byte-concat for those -// types. -// -// Writes `*entries[i].gate_up_w_out` to point at the new packed tensor. -// Marks the buffer with GGML_BACKEND_BUFFER_USAGE_WEIGHTS. The new -// context lives in `out_handles.ctx` and the backing buffer in -// `out_handles.buffer`. -// -// Returns false on alloc / size-mismatch failure; in that case -// `out_handles` is left in a state safe to free (either fully clean -// or holding a partial ctx that the caller frees on shutdown). +// Allocate a packed context + a [hidden, 2·intermediate] tensor per block +// on `backend`, then copy gate's bytes followed by up's bytes into each. +// Compatible with row-wise quants (Q4/Q5/Q6/Q8) because concat-along-dim-1 +// is byte-concat for those types. Writes `*entries[i].gate_up_w_out` to the +// new tensor and marks the buffer GGML_BACKEND_BUFFER_USAGE_WEIGHTS. +// Returns false on alloc / size-mismatch failure; `out_handles` is left in +// a state safe to free. bool pack_gate_up(ggml_backend_t backend, int hidden, int intermediate, @@ -385,13 +273,11 @@ bool pack_gate_up(ggml_backend_t backend, PackedGateUpHandles & out_handles, const char * error_tag = "causal_lm"); -// --------------------------------------------------------------------------- -// Batched greedy step loop (offline transcribe_run_batch decode) -// --------------------------------------------------------------------------- +// Batched greedy step loop (offline transcribe_run_batch decode). // The B utterances' batched step graph, as built by a family's -// build_step_graph_batched(). The field NAMES differ across families' build -// structs, so the caller fills this projection at the call site. +// build_step_graph_batched(). Field names differ across families, so the +// caller fills this projection at the call site. struct StepBatchedIO { ggml_tensor * input_ids = nullptr; // [B] i32 — token fed each step ggml_tensor * positions = nullptr; // [B] i32 — RoPE position per row @@ -420,15 +306,10 @@ struct StepLoopStats { // Run the lockstep batched greedy decode. Each row steps until it emits // `eos_id`, accumulates `max_new` generated tokens, or fills the KV window; // each emitted token is appended to generated[b]. Finished / invalid rows keep -// stepping into their own KV slab (independent memory, so a no-op for live -// rows). Polls session->poll_abort() once per step. The step graph must already -// be built and allocated on `sched`. Returns TRANSCRIBE_ERR_ABORTED on abort, +// stepping into their own KV slab (a no-op for live rows). Polls +// session->poll_abort() once per step. The step graph must already be built +// and allocated on `sched`. Returns TRANSCRIBE_ERR_ABORTED on abort, // TRANSCRIBE_ERR_GGUF on a compute failure, else TRANSCRIBE_OK. -// -// Extracted verbatim from the four causal_lm-family run_batch() drivers -// (qwen3_asr / funasr_nano / canary_qwen / granite), which were byte-identical -// here apart from granite hand-coding the f16 mask literals (0x0000 / 0xFC00 == -// ggml_fp32_to_fp16(0) / ggml_fp32_to_fp16(-inf), so this is bit-identical). transcribe_status run_batched_step_loop( transcribe_session * session, ggml_backend_sched_t sched, diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 4d007664..8f533eca 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -1,17 +1,9 @@ // src/conformer/conformer.cpp - shared Conformer encoder helpers. // -// Extracted from parakeet/encoder.cpp and cohere/encoder.cpp. The two -// files had grown to ~70% overlap, and the per-op helpers (layer_norm, -// feed_forward, the f32-friendly conv helpers, rel_shift, conv_module, -// rel_pos_mhsa, build_conformer_block, build_pre_encode) were byte- -// identical or trivially parameterizable. This module owns the -// canonical implementation; each family's encoder.cpp keeps only the -// glue that binds the shared helpers to its own weights / hparams / -// graph-driver contract. -// -// Per-family policy knobs (ConvPolicy, BlockParams) are passed in by -// the family — no environment variables are read here except by the -// shared detect_direct_pw, which is identical across families. +// See conformer.h for the API contract. Each family's encoder.cpp keeps +// only the glue binding these helpers to its weights / hparams / graph +// driver; per-family policy knobs (ConvPolicy, BlockParams) are passed in. +// No environment variables are read here except by detect_direct_pw. #include "conformer/conformer.h" #include "transcribe-log.h" @@ -25,10 +17,6 @@ namespace transcribe::conformer { -// =========================================================================== -// Low-level helpers -// =========================================================================== - ggml_tensor * named(ggml_tensor * t, const char * name) { if (t != nullptr) { ggml_set_name(t, name); @@ -37,16 +25,9 @@ ggml_tensor * named(ggml_tensor * t, const char * name) { } // LayerNorm with affine: y = gamma * (x - mean) / sqrt(var + eps) + beta. -// -// ggml_norm normalizes along ne[0], which for our [d_model, T, B] -// encoder activation is the d_model axis — exactly what LayerNorm -// wants. The affine multiply + add use ggml's broadcasting: gamma / -// beta have ne=[d_model, 1, 1, 1], which broadcasts naturally across -// the T and B axes of the activation. -// -// `gamma` is required; `beta` may be null for bias-free LN. Both -// Parakeet and Cohere ship with affine (gamma + beta) on every LN, -// but the helper supports no-bias for future families. +// ggml_norm normalizes along ne[0] (the d_model axis of [d_model, T, B]); +// gamma/beta [d_model] broadcast over T and B. `gamma` required; `beta` +// may be null for bias-free LN. ggml_tensor * layer_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * gamma, @@ -60,17 +41,9 @@ ggml_tensor * layer_norm(ggml_context * ctx, return y; } -// FeedForward: y = Linear2(SiLU(Linear1(x))) -// -// ggml_mul_mat(W, x) reads W ne[0] as the input dim, so with -// W ne=[d_in, d_out], x ne=[d_in, T, B] -> result ne=[d_out, T, B]. -// Both converters store ff*_lin1_w as ne=[d_model, d_ff] and -// ff*_lin2_w as ne=[d_ff, d_model] to match this convention. Bias -// tensors, when present, are [d_out] and broadcast along T and B. -// -// Either or both biases may be null — the optional-bias pattern is -// what lets Parakeet (bias-free) and Cohere (all biases) share the -// same code. +// FeedForward: y = Linear2(SiLU(Linear1(x))). Weights stored +// ff*_lin1_w ne=[d_model, d_ff], ff*_lin2_w ne=[d_ff, d_model]. Biases +// [d_out] broadcast along T and B; either or both may be null. ggml_tensor * feed_forward(ggml_context * ctx, ggml_tensor * x, ggml_tensor * lin1_w, @@ -90,10 +63,8 @@ ggml_tensor * feed_forward(ggml_context * ctx, return y; } -// Macaron half-residual feed-forward: x + 0.5 * FF(LN(x)). -// Used at the start (FF1) and end (FF2) of each conformer block. -// The 0.5 scale is the macaron weighting; the alternative residual -// (full 1.0) is used for the MHSA and conv sub-blocks. +// Macaron half-residual feed-forward: x + 0.5 * FF(LN(x)). The 0.5 scale +// is the macaron weighting (FF1/FF2); MHSA and conv use a full residual. ggml_tensor * macaron_ff_residual(ggml_context * ctx, ggml_tensor * x, ggml_tensor * norm_w, @@ -109,10 +80,8 @@ ggml_tensor * macaron_ff_residual(ggml_context * ctx, return ggml_add(ctx, x, y); } -// Shaw / Transformer-XL relative-position skew. Transforms the -// position scoring matrix from [pos_len, T_q, H, 1] to the same shape -// but with each row rotated so that column k holds the score for -// relative offset k. +// Shaw / Transformer-XL relative-position skew. Rotates each row of the +// [pos_len, T_q, H, 1] score matrix so column k holds relative offset k. ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x) { const int64_t pos_len = x->ne[0]; const int64_t T_q = x->ne[1]; @@ -148,11 +117,9 @@ ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x) { return y; } -// f32-friendly Conv1D (mirrors ggml_conv_1d but passes the kernel's -// real type to im2col instead of forcing f16). The vendored ggml's -// ggml_conv_1d hardcodes GGML_TYPE_F16 for the im2col output, which -// triggers a runtime assertion when fed an f32 kernel — and both the -// Parakeet and Cohere weight catalogs are fp32 for these convs. +// f32-friendly Conv1D (mirrors ggml_conv_1d but passes the kernel's real +// type to im2col instead of forcing f16). Vendored ggml's ggml_conv_1d +// hardcodes GGML_TYPE_F16 for the im2col output and asserts on an f32 kernel. ggml_tensor * conv_1d_f32(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data, @@ -176,15 +143,13 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, kernel->ne[0] * kernel->ne[1], kernel->ne[2]); // [IC*K, OC] - // F32 accumulation when the kernel is F16. The same CUDA cuBLAS - // COMPUTE_16F saturation that hits the in-block pointwise convs also - // hits this im2col + mul_mat for the subsampling conv stem when the - // converter routes those kernels to F16 (medasr / cohere / parakeet). + // F32 accumulation when the kernel is F16 (see mul_mat_f32acc in + // causal_lm.cpp for the CUDA COMPUTE_16F saturation rationale). const bool kernel_needs_f32_acc = (kernel->type == GGML_TYPE_F16); if (N == 1) { - // Single-shot path (bit-identical to the pre-batch code): flatten - // OW*N (== OW) and reshape the [OW, OC] result back to [OW, OC, 1]. + // Single-shot path: flatten OW*N (== OW) and reshape the [OW, OC] + // result back to [OW, OC, 1]. ggml_tensor * result = ggml_mul_mat(ctx, ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[1]), kernel_2d); // [OW, OC] @@ -200,9 +165,7 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, // (ow_n = b*OW + ow), so a [OW*N, OC] -> [OW, OC, N] reshape would // mislabel memory. Instead mul_mat the 3-D im2col directly: the kernel // (ne[2]=1) broadcasts across the batch, giving [OC, OW, N], then - // permute to the [OW, OC, N] = [time, channels, batch] convention. The - // per-element dot products match the N==1 path exactly (same reduction - // over IC*K), so each utterance is bit-identical to single-shot. + // permute to the [OW, OC, N] = [time, channels, batch] convention. ggml_tensor * result = ggml_mul_mat(ctx, kernel_2d, im2col); // [OC, OW, N] if (kernel_needs_f32_acc) { ggml_mul_mat_set_prec(result, GGML_PREC_F32); @@ -211,20 +174,13 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, return result; // [OW, OC, N] } -// f32-friendly 2D depthwise conv. Mirrors ggml_conv_2d_dw exactly -// but passes the kernel's real type to im2col instead of forcing -// f16. The vendored ggml's ggml_conv_2d_dw hardcodes -// GGML_TYPE_F16 for the im2col output, which then makes the -// downstream matmul try to use a `kernel_mul_mv_f32_f16_short` -// Metal kernel that isn't compiled into this ggml's Metal -// library. Forcing f32 lands the matmul on -// `kernel_mul_mv_f32_f32_short` which IS available, unblocking -// the Metal path. -// -// The alternative `ggml_conv_2d_dw_direct` hits an even harder -// wall: its CONV_2D_DW op is reported as unsupported by -// ggml_metal_op_encode_impl, so it can't run on Metal at all in -// this version. The im2col + mul_mat path here is what works. +// f32-friendly 2D depthwise conv (mirrors ggml_conv_2d_dw but passes the +// kernel's real type to im2col). CANONICAL Metal conv-quirk note: vendored +// ggml's ggml_conv_2d_dw hardcodes F16 im2col, which routes the matmul to +// a `kernel_mul_mv_f32_f16_short` Metal kernel that isn't compiled in; +// forcing f32 lands on the available `kernel_mul_mv_f32_f32_short`. The +// alternative ggml_conv_2d_dw_direct can't help: its CONV_2D_DW op is +// unsupported on Metal, so the im2col + mul_mat path here is what works. ggml_tensor * conv_2d_dw_f32(ggml_context * ctx, ggml_tensor * kernel, // [KW, KH, 1, C] ggml_tensor * data, // [W, H, C, N] @@ -283,15 +239,13 @@ ggml_tensor * conv_1d_dw_f32(ggml_context * ctx, // Offline batch: data ne=[W, C, B] with B>1. The im2col path below // collapses the batch axis (its final reshape hardcodes ne[2]=1), so // route B>1 through the direct depthwise-2D op (W, H=1, C, N=B), which - // threads the utterance batch at ne[3]. This keeps the B==1 path — - // every current caller (parakeet/cohere single-shot, AR granite) — - // byte-identical to the pre-batch code. + // threads the utterance batch at ne[3]. const int64_t B = data->ne[2]; if (B > 1) { const int64_t k = kernel->ne[0]; const int64_t C = data->ne[1]; - // ggml_conv_2d_dw_direct misbehaves for non-f32 kernels (see the - // canary_qwen note); the depthwise kernel is tiny, so cast to f32. + // ggml_conv_2d_dw_direct misbehaves for non-f32 kernels; the + // depthwise kernel is tiny, so cast to f32. ggml_tensor * knl = kernel; if (knl->type != GGML_TYPE_F32) { knl = ggml_cast(ctx, knl, GGML_TYPE_F32); @@ -355,36 +309,25 @@ ggml_tensor * add_conv_bias(ggml_context * ctx, return ggml_add(ctx, conv_out, bias_4d); } -// =========================================================================== -// Policy detection -// =========================================================================== - bool detect_direct_pw(const char * backend) { const char * env = std::getenv("TRANSCRIBE_CONV_NO_DIRECT_PW"); if (env != nullptr) return false; // user override env = std::getenv("TRANSCRIBE_CONV_DIRECT_PW"); if (env != nullptr) return true; // user override - // Vulkan: controlled A/B on AMD Renoir showed the im2col + mul_mat - // path is ~200 ms per encode faster than the direct mul_mat path - // for f32 weights, so keep im2col as the default. Metal and CPU - // benefit from the direct path. + // Vulkan defaults to im2col: on AMD Renoir, im2col + mul_mat measured + // ~200 ms/encode faster than direct for f32 weights. Metal/CPU prefer + // direct. if (backend != nullptr && std::strstr(backend, "Vulkan") != nullptr) { return false; } return true; } -// =========================================================================== -// Conv module (pointwise -> GLU -> depthwise -> BN -> SiLU -> pointwise) -// =========================================================================== - -// Operates on the post-LayerNorm activation; the LN is applied OUTSIDE -// this helper by the block forward. By default, pointwise convs (k=1) -// are direct mul_mats in [d_model, T] layout, avoiding im2col overhead. -// Set TRANSCRIBE_CONV_NO_DIRECT_PW=1 to fall back to the im2col path. -// BatchNorm uses precomputed fused scale + bias (2 ops instead of 4); -// LayerNorm uses an unfused on-the-fly per-channel normalize plus -// affine, with weights from BlockView::conv_ln_*. +// Conv module (pointwise -> GLU -> depthwise -> BN/LN -> SiLU -> pointwise). +// Operates on the post-LayerNorm activation (LN applied by the caller). +// Pointwise convs (k=1) are direct mul_mats in [d_model, T] layout by +// default; TRANSCRIBE_CONV_NO_DIRECT_PW=1 forces the im2col path. BatchNorm +// uses fused scale + bias; LayerNorm normalizes per-channel on the fly. ggml_tensor * conv_module(ggml_context * ctx, ggml_tensor * x, const BlockView & b, @@ -405,21 +348,14 @@ ggml_tensor * conv_module(ggml_context * ctx, const int64_t d_model = x->ne[0]; // Utterance batch lives at ne[2] of the [d_model, T, B] activation. - // B == 1 for single-shot (parakeet / cohere today); the offline - // batched encoder passes B > 1. Every view/reshape below threads B so - // the B == 1 path is bit-identical to the pre-batch code. + // B == 1 for single-shot; the offline batched encoder passes B > 1. const int64_t B = x->ne[2]; if (policy.direct_pw) { // Pointwise conv 1 as direct mul_mat in [d_model, T, B] layout. // Kernel ne=[1, d_model, 2*d_model] → reshape to [d_model, 2*d_model]. - // F32 accumulation when the weight is F16: on CUDA, F16 weights take - // the cuBLAS COMPUTE_16F path whose accumulator saturates at ~6.5e4. - // Families with large activation magnitudes (e.g. medasr's scaled - // residuals push intermediates to ~2e6) overflow to NaN. set_prec - // forces F32 accumulation; no-op on CPU/Metal, slight perf cost on - // CUDA. Skip for non-F16 weights (BF16 already COMPUTE_32F; quant - // dispatches MMQ). + // Force F32 accumulation for F16 weights (see mul_mat_f32acc in + // 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); @@ -576,15 +512,10 @@ ggml_tensor * conv_module(ggml_context * ctx, x = ggml_add(ctx, x, bias_r); } - // Post-depthwise normalisation: BN (fused mul + add) or LN - // (per-channel mean/std normalize + affine). Both produce the same - // [T, d_model] shape; the SiLU below is identical. - // - // At this point x has ne = [T, d_model, 1, 1]. NeMo's LayerNorm - // normalises over the feature axis (d_model), matching the - // `layer_norm` helper exactly when applied to a [T, d_model] - // tensor — but `layer_norm` reads ne[0] as the feature axis, so we - // permute to [d_model, T] first, apply, and permute back. + // Post-depthwise normalisation: BN (fused mul + add) or LN (per-channel + // mean/std + affine). x is [T, d_model, 1, 1] here; `layer_norm` reads + // ne[0] as the feature axis, so for LN we permute to [d_model, T], apply, + // and permute back. if (params.conv_norm_type == BlockParams::ConvNormType::LayerNorm) { x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); x = layer_norm(ctx, x, b.conv_ln_w, b.conv_ln_b); @@ -594,7 +525,6 @@ ggml_tensor * conv_module(ggml_context * ctx, b.conv_bn_fused_scale, b.conv_bn_fused_bias); } - // SiLU activation. x = ggml_silu(ctx, x); if (policy.direct_pw) { @@ -626,10 +556,8 @@ ggml_tensor * conv_module(ggml_context * ctx, return x; } -// =========================================================================== -// Relative-position multi-head self-attention -// =========================================================================== - +// Relative-position multi-head self-attention. +// // use_flash = true -> ggml_flash_attn_ext (fused GPU kernel, needs // a dk template in the Metal backend) // use_flash = false -> manual mul_mat + soft_max_ext + mul_mat (works @@ -664,12 +592,10 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, const int att_context_right = params.att_context_right; const int head_dim = d_model / n_head; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); - // Query/key-value split (see the x_q / k_full contracts in - // conformer.h). kv_cached: K/V arrive pre-projected and x is the - // query-only activation. q_sliced: queries from x_q, K/V projected - // from x. rect covers both — the score geometry is rectangular - // [T_kv, T_q]. Neither set keeps T_q == T_kv and the historical - // topology bit-for-bit. + // Query/key-value split (see the x_q / k_full contracts in conformer.h). + // kv_cached: K/V arrive pre-projected, x is the query-only activation. + // q_sliced: queries from x_q, K/V projected from x. rect covers both — + // the score geometry is rectangular [T_kv, T_q]. const bool kv_cached = (k_full != nullptr && v_full != nullptr); const bool q_sliced = !kv_cached && (x_q != nullptr && x_q != x); const bool rect = kv_cached || q_sliced; @@ -687,25 +613,13 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, (long long)pos_len, (long long)(T_q + T_kv - 1)); return nullptr; } - // Utterance batch at ne[2] of the [d_model, T, B] activation. After the - // head split below the batch moves to ne[3] (heads take ne[2]). When - // batched we force the manual attention path: ggml_flash_attn_ext's - // additive mask (matrix_bd, which is content-dependent and therefore - // differs per utterance) is not validated for a per-batch ne[3] here, - // whereas the manual mul_mat + soft_max path broadcasts over the batch - // cleanly. Single-shot (B == 1) keeps its existing flash path exactly. + // Utterance batch at ne[2] (moves to ne[3] after the head split). Flash + // works batched (ggml_flash_attn_ext accepts the per-utterance rel-pos + // mask at ne[3] == B). Rectangular geometry forces manual: flash requires + // the mask's ne[1] padded to GGML_KQ_MASK_PAD, which the [T_kv, T_q] + // streaming mask does not satisfy. (TRANSCRIBE_NO_FLASH=1 also forces + // manual, for the bit-exact CPU tensor gate.) const int64_t B = x->ne[2]; - // Flash attention works batched: ggml_flash_attn_ext accepts the - // per-utterance rel-pos mask at ne[3] == B, and the batched encoder - // output is bit-identical to single-shot flash (verified on CPU and - // Metal). So the batch axis does not change the attention path. - // TRANSCRIBE_NO_FLASH=1 forces the manual mul_mat + soft_max path for - // both single-shot and batched (used by the bit-exact CPU tensor gate, - // since flash casts the rel-pos mask to F16 while manual stays F32). - // Rectangular geometry also forces manual: ggml_flash_attn_ext - // requires the mask's ne[1] padded to GGML_KQ_MASK_PAD, which the - // rectangular [T_kv, T_q] streaming mask does not satisfy — and at - // streaming geometry the manual path is faster anyway on CPU. const bool flash = use_flash && !rect; // Local-attention bookkeeping. With both window sides non-negative @@ -720,9 +634,8 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, const int W_left = is_local ? att_context_left : 0; const int W_right = is_local ? att_context_right : 0; - // ----- Q, K, V, P projections --------------------------------- - // Q, K, V may have bias (Cohere) or not (Parakeet). Pos projection - // (attn_pos_w) never has a bias in either family. + // Q, K, V, P projections. Q/K/V may have bias (Cohere) or not + // (Parakeet); the pos projection (attn_pos_w) never has a bias. ggml_tensor * q = ggml_mul_mat(ctx, b.attn_q_w, q_sliced ? x_q : x); if (b.attn_q_b != nullptr) q = ggml_add(ctx, q, b.attn_q_b); // KV-cache mode: keys/values arrive pre-projected (bias included). @@ -738,10 +651,8 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, ? ggml_mul_mat(ctx, b.attn_pos_w, pos_emb) : nullptr; - // ----- Split heads -------------------------------------------- - // pos_bias_u/v broadcast onto [head_dim, n_head, T, 1] BEFORE - // the permute that moves T past n_head. - // Head split carries the batch axis at ne[3]: [head_dim, n_head, T, B]. + // Split heads: [head_dim, n_head, T, B] (batch at ne[3]). pos_bias_u/v + // broadcast onto this BEFORE the permute that moves T past n_head. q = ggml_reshape_4d(ctx, q, head_dim, n_head, T_q, B); ggml_tensor * q_u = ggml_add(ctx, q, b.attn_pos_u); ggml_tensor * q_v = ggml_add(ctx, q, b.attn_pos_v); @@ -770,8 +681,7 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, p = pos_proj; } - // ----- Position mask / bias ----------------------------------- - // matrix_bd = rel_shift(q_v @ p^T), truncated to [T_q, T_q]. + // Position mask / bias: matrix_bd = rel_shift(q_v @ p^T), truncated. ggml_tensor * matrix_bd = ggml_mul_mat(ctx, p, q_v); // Local-attention pad/slice. The standard rel_shift trick assumes @@ -897,32 +807,26 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, // [T_q, T_q, n_head, 1] attention scores. ggml_tensor * kq = ggml_mul_mat(ctx, k, q_u); - // Add relative position bias, then scale inside soft_max_ext. + // Add relative position bias; soft_max_ext fuses the scale. kq = ggml_add(ctx, kq, matrix_bd); - - // ggml_soft_max_ext fuses the scale into softmax. ggml_tensor * kq_soft = ggml_soft_max_ext(ctx, kq, /*mask=*/nullptr, scale, /*max_bias=*/0.0f); - // V needs to be transposed for the value matmul. - // V is [head_dim, T_q, n_head, 1]. We need [T_q, head_dim, n_head, 1]. + // V transposed for the value matmul: [T_q, head_dim, n_head, 1]. ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v, 1, 0, 2, 3)); // attn_weights @ V: [head_dim, T_q, n_head, 1] o = ggml_mul_mat(ctx, v_t, kq_soft); - // Merge heads: permute [head_dim, T_q, n_head] -> [head_dim, - // n_head, T_q] then reshape to [d_model, T_q]. The permute - // makes heads contiguous for the reshape. + // Merge heads: permute -> [head_dim, n_head, T_q] (contiguous for + // the reshape) then reshape to [d_model, T_q]. o = ggml_permute(ctx, o, 0, 2, 1, 3); o = ggml_cont(ctx, o); } // Collapse heads back to d_model, keeping the batch at ne[2]. - // reshape_3d with B == 1 is equivalent to the prior reshape_2d. o = ggml_reshape_3d(ctx, o, d_model, T_q, B); - // ----- Output linear ----------------------------------------- o = ggml_mul_mat(ctx, b.attn_out_w, o); if (b.attn_out_b != nullptr) o = ggml_add(ctx, o, b.attn_out_b); return o; @@ -939,9 +843,7 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, const BlockParams & params, const BlockObserver * obs) { - // Observer dispatch helper. Inlines to nothing when `obs` is null - // or its callback is null, so blocks without instrumentation pay - // zero overhead. Tags are compile-time string literals. + // Observer dispatch helper (no-op when obs / its callback is null). const auto notify = [&](const char * tag, ggml_tensor * t) { if (obs != nullptr && obs->on_point != nullptr) { obs->on_point(obs->user, tag, t); @@ -958,23 +860,15 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, // Self-attention with relative position. Full residual. // // Streaming carry (params.streaming_channel_in != nullptr): - // 1. Compute the new cache as - // new_cache = concat(prev_cache[T_q_new:], x_norm) - // and cpy it into streaming_channel_out (per NeMo: - // q_keep_size = T_q_new with cache_drop_size = 0; the cache - // stays at size T_cache). KV-cache mode rotates last_k/last_v - // instead (see kv_mode below). - // 2. Virtualize x_norm for the K/V side: - // x_norm_virtual = concat(prev_cache, x_norm) - // so rel_pos_mhsa attends over T_virtual = T_cache + T_q_new - // keys. Queries are sliced to the T_q_new new frames (x_q_new), - // so the attention output already has T_q_new rows — no output - // slice needed. The pos_emb / chunked_mask are sized for that + // 1. new_cache = concat(prev_cache[T_q_new:], x_norm), cpy'd into + // streaming_channel_out (NeMo q_keep_size = T_q_new, cache stays at + // T_cache). KV-cache mode rotates last_k/last_v instead (kv_mode). + // 2. Virtualize x_norm for K/V: concat(prev_cache, x_norm) so + // rel_pos_mhsa attends over T_virtual = T_cache + T_q_new keys. + // Queries are sliced to T_q_new (x_q_new), so the output already + // has T_q_new rows; pos_emb / chunked_mask are sized for the // rectangular [T_virtual, T_q_new] geometry. - // - // T_q_new is taken from the running x's ne[1] (x is the post-pre- - // encode + post-FF1 running tensor; its time dim hasn't changed - // yet at this point). + // T_q_new comes from the running x's ne[1]. { const bool streaming = (params.streaming_channel_in != nullptr); // KV-cache streaming: keys/values are cached pre-projected, so @@ -1015,16 +909,10 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, }; if (streaming && !kv_mode) { - // Emit the new cache slot before we virtualize x_norm - // for attention. Source: prev_cache tail (T_cache - T_q_new - // rows) concatenated with this chunk's x_norm (T_q_new - // rows) — total T_cache rows along ne[1]. - // - // For T_q_new < T_cache (the common case; nemotron has - // T_q_new=12, T_cache=70), prev_cache_tail is a [d_model, - // T_cache - T_q_new] view. For T_q_new >= T_cache the - // new cache is just the last T_cache rows of x_norm — not - // exercised today but the guard is cheap. + // Emit the new cache slot (T_cache rows) before virtualizing + // x_norm: prev_cache tail (T_cache - T_q_new rows) concatenated + // with this chunk's x_norm, or x_norm's last T_cache rows when + // T_q_new >= T_cache. if (params.streaming_channel_out != nullptr && params.streaming_graph != nullptr) { @@ -1040,8 +928,7 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, new_cache = ggml_concat(ctx, prev_tail, x_norm, /*dim=*/1); } else { - // x_norm holds at least T_cache rows — take its - // last T_cache. + // x_norm's last T_cache rows. new_cache = ggml_view_2d( ctx, x_norm, x_norm->ne[0], T_cache, @@ -1117,21 +1004,18 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, b.ff2_lin2_w, b.ff2_lin2_b); notify("after_ff2", x); - // Final per-block LayerNorm (the block's output transform). + // Final per-block LayerNorm. x = layer_norm(ctx, x, b.norm_out_w, b.norm_out_b); notify("out", x); return x; } -// =========================================================================== -// Pre-encode (DwStridingSubsampling) -// =========================================================================== +// Pre-encode (DwStridingSubsampling). namespace { -// Name a tensor as `.` if prefix is non-null. Returns -// t unchanged. Used below to reproduce Parakeet's dump-point naming -// inside the shared helper without branching on family. +// Name a tensor as `.` if prefix is non-null. Returns t +// unchanged. ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, const char * suffix) { @@ -1144,8 +1028,8 @@ ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, } // namespace -// Pre-encode subsampling stack. See conformer.py:206-328 -// (DwStridingSubsampling). The op order matches NeMo: +// Pre-encode subsampling stack. Op order matches NeMo's +// DwStridingSubsampling (conformer.py:206-328): // // conv0 (standard, 1->C, k=3 s=2 p=1) -> bias -> ReLU // conv2 (depthwise, C->C, groups=C) -> bias -> conv3 (pointwise, C->C) -> bias -> ReLU @@ -1154,16 +1038,8 @@ ggml_tensor * name_prefixed(ggml_tensor * t, const char * prefix, // feature axis matching `pre_encode_in = channels * (n_mels / 8)` // linear (out_w, out_b) -> [d_model, T_enc, 1, 1] // -// In NeMo's torch.nn.Sequential, the bias is fused into each conv. -// We replicate that here with an explicit ggml_add per conv -// (add_conv_bias is a no-op when the bias slot is null). -// -// The 2-D depthwise convs (conv2, conv5) honor -// policy.direct_dw_in_pre_encode. Parakeet keeps this false (uses the -// f32-friendly im2col helper on every backend, including ones where -// the direct path would work); Cohere follows its detect_direct_dw -// result. See the policy comment in conformer.h for why the two -// families historically disagree on this. +// Each conv's fused bias is an explicit ggml_add (no-op when null). The 2-D +// depthwise convs (conv2, conv5) honor policy.direct_dw_in_pre_encode. ggml_tensor * build_pre_encode(ggml_context * ctx, const PreEncodeView & pe, ggml_tensor * mel_in, @@ -1172,11 +1048,10 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, const char * error_tag, PreEncodeValidMasks * valid_masks) { - // Utterance batch and a helper that, in the masked-subsampling path, - // zeros each conv intermediate's padded time region after a ReLU. The - // mask is a graph input ne=[1, H_stage, 1, B] broadcasting over freq - // (ne[0]) and channels (ne[2]); the driver fills it per utterance. `slot` - // receives the allocated handle. No-op when valid_masks is null. + // Masked-subsampling helper: zeros each conv intermediate's padded time + // region after a ReLU. The mask is a graph input ne=[1, H_stage, 1, B] + // (broadcasts over freq and channels); `slot` receives the handle for + // the driver to fill. No-op when valid_masks is null. const int64_t pe_batch = mel_in->ne[3]; auto apply_valid_mask = [&](ggml_tensor * x, ggml_tensor ** slot, const char * mname) -> ggml_tensor * { @@ -1189,34 +1064,24 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, return ggml_mul(ctx, x, m); }; - // Transpose the mel input from its natural [T_mel, n_mels, 1, 1] - // layout (matching the C++ row-major MelFrontend buffer) to - // [n_mels, T_mel, 1, 1] = [W=F, H=T, IC, N] which is what - // ggml_conv_2d expects under the NHWC `[B, H=T, W=F, C]` - // convention. The conv kernel is not spatially symmetric, so the - // axis order matters — if F and T are swapped the conv math - // diverges. ggml_permute returns a non-contiguous view; the - // ggml_cont materializes it for the conv. + // Transpose the mel input from [T_mel, n_mels, 1, 1] to + // [n_mels, T_mel, 1, 1] = [W=F, H=T, IC, N], the NHWC order + // ggml_conv_2d expects. The kernel is not spatially symmetric, so F/T + // order matters — swapping them diverges the conv math. ggml_tensor * x = ggml_permute(ctx, mel_in, /*axis0=*/1, /*axis1=*/0, /*axis2=*/2, /*axis3=*/3); x = ggml_cont(ctx, x); x = name_prefixed(x, name_prefix, "mel_t"); - // Causal pre_encode: NeMo's CausalConv2D applies F.pad(left=k-1, - // right=stride-1, left=k-1, right=stride-1) on both spatial axes - // before the underlying conv (kernel=3, stride=2, padding=0). We - // replicate that with explicit zero pads on dim 0 (W = freq) and - // dim 1 (H = time), then call the conv with p=0. The op-side - // (k-1)/2 symmetric padding path is what every offline variant - // takes. + // 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) { if (!causal_pe) return t; - // For k=3 / s=2: left=2, right=1 on both axes. Pad as zero - // F32 leaves; ggml_fill writes the constant after compute- - // buffer alloc, identical to how rel_pos_mhsa zero-pads. + // For k=3 / s=2: left=2, right=1 on both axes (zero F32 pads). const int left = 2, right = 1; auto make_pad = [&](int width_w, int width_h) { ggml_tensor * p = ggml_new_tensor_4d( @@ -1256,9 +1121,8 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, 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) - // See conv_2d_dw_f32 above for the Metal backstory on why the - // im2col path is used when direct_dw_in_pre_encode is false. + // 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, @@ -1314,34 +1178,21 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, "pre_encode.valid_mask.s3"); // At this point ne = [F'=16, T_enc, channels=256, 1] where - // T_enc = floor(T_mel / 8) (with the per-conv floor formula). - // We need to flatten (F', C) into a single feature axis matching - // pre_encode_in = channels * (n_mels / subsampling_factor). - // - // The reference flattens with C as the slower sub-axis and F' - // as the faster, giving flat index = c*F' + f'. In ggml - // fast-to-slow ne, that's the layout you get by permuting to put - // F' on axis 0 (already there) and C on axis 1 (currently axis - // 2), then collapsing axes 0 and 1. - // + // T_enc = floor(T_mel / 8). Flatten (F', C) into one feature axis + // matching pre_encode_in = channels * (n_mels / subsampling_factor). + // The reference uses flat index = c*F' + f' (C slow, F' fast), which in + // ggml ne is: permute F' to axis 0, C to axis 1, then collapse 0 and 1: // permute (0, 2, 1, 3): [F', T, C, 1] -> [F', C, T, 1] - // ggml_cont: (reshape requires a contiguous tensor) - // reshape_2d to [F'*C, T, 1, 1] = [pre_encode_in, T_enc, 1, 1] + // ggml_cont; reshape -> [F'*C, T] = [pre_encode_in, T_enc] const int64_t F_prime = x->ne[0]; const int64_t T_enc = x->ne[1]; const int64_t C = x->ne[2]; - // Utterance batch rides the conv "N" axis (ne[3]) through the whole - // stem; after the flatten + linear below it lands on the activation's - // ne[2] == B convention the conformer blocks expect. B == 1 for - // single-shot, so the 3-D reshape collapses to the prior 2-D shape. + // Utterance batch rides the conv "N" axis (ne[3]) through the stem; the + // flatten + linear below land it on ne[2] == B for the conformer blocks. const int64_t B = x->ne[3]; const int64_t pre_encode_in = F_prime * C; - // Sanity check against the catalog: the linear layer expects - // exactly this many input features. If the math drifts (e.g. a - // future converter changes the subsampling layout), surface it - // here as a clear assertion rather than letting mul_mat fail - // with a confusing shape mismatch. + // Sanity check: the linear layer expects exactly pre_encode_in features. if (pe.out_w != nullptr && pre_encode_in != pe.out_w->ne[0]) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "%s encoder: pre_encode_in mismatch: " @@ -1361,16 +1212,11 @@ ggml_tensor * build_pre_encode(ggml_context * ctx, x = ggml_reshape_3d(ctx, x, pre_encode_in, T_enc, B); x = name_prefixed(x, name_prefix, "flat"); - // Linear projection to d_model. ggml_mul_mat(W, x): - // W ne=[in, out], x ne=[in, T] -> out ne=[out, T] - // The catalog stores out_w as ne=[pre_encode_in, d_model] which - // matches this convention exactly. + // Linear projection to d_model (out_w ne=[pre_encode_in, d_model]). x = ggml_mul_mat(ctx, pe.out_w, x); x = name_prefixed(x, name_prefix, "linear"); - // Add bias [d_model] -> broadcast to [d_model, T_enc, 1, 1]. - // Same trick as the conv biases above: reshape to a 4-D view - // with the channel axis in position 0 and size-1 dims elsewhere. + // Add bias [d_model], broadcast to [d_model, T_enc, 1, 1]. if (pe.out_b != nullptr) { const int64_t d_model = pe.out_b->ne[0]; ggml_tensor * bias_4d = ggml_reshape_4d(ctx, pe.out_b, diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index db4d9bfa..59c80b69 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -1,57 +1,21 @@ // src/conformer/conformer.h - shared Conformer encoder helpers. // -// Both Parakeet and Cohere-ASR (and any future family built on a -// NeMo-style Conformer encoder) share the same block topology: +// Family-agnostic helpers for the NeMo-style Conformer encoder shared by +// Parakeet, Cohere-ASR, and future variants. Block topology: // -// pre_encode (DwStridingSubsampling: 1 standard conv + 2 depthwise -// -separable conv pairs, ReLU-interleaved, then a linear -// projection to d_model) -// N x conformer block { -// macaron FF1 (0.5x residual) -// rel-pos MHSA (1.0x residual) -// conv module (pointwise-GLU -> depthwise -> BN -> SiLU -> pointwise) -// macaron FF2 (0.5x residual) -// final per-block LayerNorm -// } +// pre_encode (DwStridingSubsampling) -> +// N x { macaron FF1 (0.5x) -> rel-pos MHSA -> conv module +// (pw-GLU -> depthwise -> BN -> SiLU -> pw) -> macaron FF2 (0.5x) +// -> per-block LayerNorm } // -// This header exposes that machinery as a set of family-agnostic -// helpers driven by three structs: -// -// - PreEncodeView / BlockView : nullable-pointer projections over -// whatever per-family weight struct the family's weights.h -// defines. Families with biases (Cohere) fill the bias slots; -// bias-free families (Parakeet) leave them nullptr. Every helper -// that takes a bias pointer treats nullptr as "no bias" and skips -// the add — this is the optional-bias pattern. -// -// - ConvPolicy : per-family dispatch choice for the pointwise and -// depthwise convs. The pointwise policy is identical across -// families (detect_direct_pw lives here, shared) but the -// depthwise policy splits between "inside the conformer block" -// and "inside pre_encode" because the two sites have different -// tensor shapes and historical backend behavior. Each family -// populates its own ConvPolicy; the shared helpers never read -// environment variables directly for dw dispatch. -// -// - BlockParams : the scalar knobs the block forward needs -// (d_model, n_head, conv_kernel, kv_type, use_flash, policy). -// -// Sub-helpers (macaron_ff_residual, rel_pos_mhsa, conv_module) are -// also exposed at namespace scope so a family that hand-builds one or -// more blocks for debug-dump purposes (currently Parakeet's block 0) -// can drive the same code paths without going through -// build_conformer_block. No helper in this module ever calls -// ggml_set_name or ggml_set_output on intermediates — naming and -// output marking are the family's responsibility, and the one helper -// that does name intermediates (build_pre_encode) takes an explicit -// name prefix. -// -// The f32-friendly conv helpers (conv_1d_f32, conv_2d_dw_f32, -// conv_1d_dw_f32) exist to work around ggml kernel limitations on -// Metal — see the block comment above conv_2d_dw_f32 in conformer.cpp -// for the full story. Callers should use these, NOT ggml's built-in -// ggml_conv_1d / ggml_conv_2d_dw, whenever the kernel weights are -// f32 and the compute must run on Metal. +// Driven by three structs: PreEncodeView / BlockView (nullable-pointer +// projections; null bias = skip the add), ConvPolicy (per-family pointwise/ +// depthwise conv dispatch), and BlockParams (scalar knobs). Sub-helpers are +// exposed at namespace scope for families that hand-build a block for +// debug-dump. Helpers name no intermediates except build_pre_encode (explicit +// prefix). The f32-friendly conv helpers (conv_1d_f32, conv_2d_dw_f32, +// conv_1d_dw_f32) work around ggml kernel limitations on Metal (see +// conv_2d_dw_f32 in conformer.cpp); use them for f32 kernels on Metal. #pragma once @@ -59,23 +23,12 @@ namespace transcribe::conformer { -// =========================================================================== -// Constants -// =========================================================================== - -// LayerNorm + BatchNorm epsilon. Both default to 1e-5 in NeMo and -// are NOT stored in the GGUF. Hardcoded here to match. If a future -// variant changes either value, plumb it through BlockParams rather -// than bumping this constant. +// LayerNorm + BatchNorm epsilon. Both default to 1e-5 in NeMo and are NOT +// stored in the GGUF, so hardcode to match. constexpr float kLayerNormEps = 1e-5f; -// =========================================================================== -// Views over per-family weight structs -// =========================================================================== - // Nullable-pointer projection over the pre_encode weights. The family -// (parakeet / cohere) fills this from its own weights struct at the -// call site — see to_view() helpers in each encoder.cpp. +// fills this from its own weights struct at the call site. struct PreEncodeView { ggml_tensor * conv0_w = nullptr; // [KW=3, KH=3, IC=1, OC=channels] ggml_tensor * conv0_b = nullptr; // [channels] @@ -93,9 +46,7 @@ struct PreEncodeView { }; // Nullable-pointer projection over one Conformer block's weights. -// Every `_b` (bias) slot is genuinely optional — leave nullptr for -// families that ship without that bias. The matching helper skips the -// add when the pointer is null. +// Every `_b` (bias) slot is optional — null = skip the add. struct BlockView { // Macaron FF1. ggml_tensor * norm_ff1_w = nullptr; // [d_model] @@ -129,14 +80,9 @@ struct BlockView { ggml_tensor * conv_dw_b = nullptr; // [d_model], nullable ggml_tensor * conv_pw2_w = nullptr; // [1, d_model, d_model] ggml_tensor * conv_pw2_b = nullptr; // [d_model], nullable - // Post-depthwise normalisation. Exactly one pair is populated per - // block, depending on BlockParams::conv_norm_type. - // BatchNorm path: conv_bn_fused_scale / conv_bn_fused_bias - // (precomputed at load from raw BN weight/bias + - // running_mean/var). Offline Conformer variants. - // LayerNorm path: conv_ln_w / conv_ln_b (raw LN scale/bias; the - // per-channel mean/std is computed at inference). - // Streaming variants (nemotron-speech-streaming). + // Post-depthwise normalisation. Exactly one pair is populated per block + // per BlockParams::conv_norm_type: BatchNorm (conv_bn_fused_*, fused at + // load) or LayerNorm (conv_ln_*, per-channel mean/std at inference). ggml_tensor * conv_bn_fused_scale = nullptr; // [d_model] ggml_tensor * conv_bn_fused_bias = nullptr; // [d_model] ggml_tensor * conv_ln_w = nullptr; // [d_model] @@ -155,69 +101,33 @@ struct BlockView { ggml_tensor * norm_out_b = nullptr; }; -// =========================================================================== -// Policy and params -// =========================================================================== - -// Per-family conv dispatch policy. direct_pw is shared -// (detect_direct_pw produces the same answer for every family today) -// but direct_dw splits between block-internal and pre_encode sites: -// -// - direct_dw_in_block : the conformer block's 1-D depthwise -// conv (after GLU). Parakeet takes the -// direct path on every backend; Cohere -// takes it only on Vulkan/CUDA and uses -// im2col on Metal/CPU. See each family's -// detect_direct_dw() for the exact -// policy. -// -// - direct_dw_in_pre_encode : the pre_encode 2-D depthwise conv -// (stride-2 downsample). Historically -// Parakeet hardcoded the im2col path here -// (direct_dw_in_pre_encode = false) -// regardless of the in-block flag; -// Cohere routes both sites through the -// same detect_direct_dw value. Splitting -// the two lets each family keep its -// historical behavior exactly, no quiet -// unification. -// Defaults are deliberately conservative: direct_pw is true because -// pointwise convs are direct mul_mat on every backend, but both -// direct_dw_* default to false (im2col path). A future family that -// default-initializes ConvPolicy{} and forgets to set these will get -// the im2col path — slower than direct on Vulkan but safe on Metal, -// where the direct 2D depthwise kernel is not implemented for all -// shapes and the im2col + mul_mat path is the reliable fallback. -// Parakeet and Cohere populate all three fields explicitly at the -// call site in build_encoder_graph, so the defaults only matter for -// new callers. +// Per-family conv dispatch policy. direct_pw is shared (detect_direct_pw is +// the same for every family today); direct_dw splits between the block site +// (direct_dw_in_block: the conformer block's 1-D depthwise after GLU) and the +// pre_encode site (direct_dw_in_pre_encode: the stride-2 2-D depthwise), since +// the two have different shapes and per-family backend choices. Defaults are +// conservative: direct_pw true (pointwise is direct mul_mat everywhere), both +// direct_dw_* false (im2col), which is safe on Metal where the direct 2-D +// depthwise kernel is not implemented for all shapes. struct ConvPolicy { bool direct_pw = true; bool direct_dw_in_block = false; bool direct_dw_in_pre_encode = false; - // Causal pre_encode convolutions. NeMo's cache-aware streaming - // (`is_causal=true`) swaps every Conv2d in ConvSubsampling for - // CausalConv2D, which applies `F.pad(left=k-1, right=stride-1)` on - // both spatial axes before calling the underlying conv with p=0. - // For k=3 / s=2 that is (left=2, right=1) per axis — different - // total padding from the offline (k-1)/2 symmetric path, and - // shifts both the freq output dim (e.g. 128 → 65 → 33 → 17 - // instead of 64 → 32 → 16) and the time output dim. False on every - // offline Parakeet variant; true on nemotron-speech-streaming-en. + // 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 + // is (left=2, right=1), different from the offline (k-1)/2 symmetric + // path, and shifts both the freq and time output dims. False on every + // offline variant; true on nemotron-speech-streaming-en. bool causal_pre_encode = false; }; -// Pointwise conv dispatch auto-detect. Identical policy across -// families today (Parakeet and Cohere return the same answer for -// every backend) so this lives here rather than in each encoder.cpp. +// Pointwise conv dispatch auto-detect (same answer for every family today). // Env overrides: TRANSCRIBE_CONV_NO_DIRECT_PW=1 forces im2col, -// TRANSCRIBE_CONV_DIRECT_PW=1 forces direct. -// -// Policy note (Vulkan): on AMD Renoir, a controlled A/B showed the -// im2col + mul_mat path is ~200 ms per encode faster than the direct -// mul_mat path for f32 weights, so Vulkan defaults to im2col. Metal -// and CPU prefer direct. +// TRANSCRIBE_CONV_DIRECT_PW=1 forces direct. Vulkan defaults to im2col: on +// AMD Renoir, im2col + mul_mat measured ~200 ms/encode faster than direct for +// f32 weights. Metal and CPU prefer direct. bool detect_direct_pw(const char * backend); // Scalar knobs for the block forward. All fields required except the @@ -277,9 +187,12 @@ struct BlockParams { int conv_context_left = -1; int conv_context_right = -1; - // Optional post-GLU valid-frame mask for the conv module. Shape - // [T, 1, 1, 1], values 1.0 for valid frames and 0.0 for padded - // overhang. NeMo applies pad_mask after pointwise_conv1+GLU and + // Optional post-GLU valid-frame mask for the conv module. Values 1.0 + // for valid frames and 0.0 for padded overhang; multiplied onto x at + // ne=[T, d_model, B, 1], so it broadcasts over d_model (ne[1]). Shape + // is [T, 1, 1, 1] for the single-utterance / buffered-streaming case + // and [T, 1, B, 1] for variable-length offline batching (the B axis + // lands in ne[2]). NeMo applies pad_mask after pointwise_conv1+GLU and // before the depthwise convolution; null keeps the historical path. ggml_tensor * conv_pad_mask = nullptr; @@ -290,90 +203,49 @@ struct BlockParams { enum class ConvNormType { BatchNorm, LayerNorm }; ConvNormType conv_norm_type = ConvNormType::BatchNorm; - // ---- Streaming (cache-aware) inputs/outputs ---- - // - // All four pointers are nullptr in offline mode and the block - // runs the existing graph topology. When non-null, the block - // runs the streaming path: - // - // streaming_channel_in per-layer cache_last_channel tensor - // from the previous chunk (persistent backend buffer). Shape - // [d_model, T_cache, 1, 1]. Concat-prepended onto the post- - // attn-LN x before Q/K/V projection ("virtual T" approach): - // the block runs attention on T_virtual = T_cache + T_q_new - // positions and slices the output to the last T_q_new rows. - // - // streaming_channel_out output tensor (allocated by the caller - // in the per-call compute_ctx) that rel_pos_mhsa fills via - // ggml_cpy with the tail of x_norm (size T_q_new, the new - // cache slot per NeMo's q_keep_size = T_new rule with - // cache_drop_size = 0). After graph_compute, the driver - // rotates this into the persistent cache buffer. - // - // streaming_time_in per-layer cache_last_time tensor from - // the previous chunk. Shape [k_minus_1, d_model, 1, 1]. - // Replaces the zero left-pad in the depthwise conv (the - // causal pad on streaming variants is (k-1, 0); the k-1 left - // frames are now the previous chunk's post-pw1+GLU tail). - // - // streaming_time_out output tensor (compute_ctx) that - // conv_module fills with the last k_minus_1 frames of the - // post-pw1+GLU input (which becomes the next chunk's left - // pad). - // - // The pos_emb tensor (passed separately to rel_pos_mhsa) and the - // attn_chunked_mask (above) are sized for the virtual T when - // streaming, not the offline T_enc. The caller is responsible for - // building them at the correct size. + // Streaming (cache-aware) inputs/outputs. All null in offline mode + // (existing graph topology). When non-null the block runs the streaming + // path with a "virtual T" attention window (T_cache + T_q_new): + // streaming_channel_in [d_model, T_cache] cache_last_channel from the + // previous chunk, concat-prepended before Q/K/V projection. + // streaming_channel_out new cache slot (tail of x_norm, size T_q_new); + // the driver rotates it into the persistent buffer after compute. + // streaming_time_in [k-1, d_model] cache_last_time replacing the + // depthwise conv's zero left-pad (causal pad is (k-1, 0)). + // streaming_time_out last k-1 post-pw1+GLU frames -> next left pad. + // pos_emb and attn_chunked_mask must be sized for the virtual T. ggml_tensor * streaming_channel_in = nullptr; ggml_tensor * streaming_channel_out = nullptr; ggml_tensor * streaming_time_in = nullptr; ggml_tensor * streaming_time_out = nullptr; - // When streaming, the number of new (post-pre-encode) encoder - // frames being produced this call. Used to: - // - slice rel_pos_mhsa's attention output back to the new rows - // - decide the source slice for streaming_channel_out / _time_out - // Ignored in offline mode (streaming_channel_in == nullptr). + // When streaming, the number of new encoder frames produced this call + // (slices the attention output and the cache-out source). Offline: 0. int streaming_T_q_new = 0; - // Optional streaming KV cache for this block (all four set - // together, or none): persistent [d_model, T_cache] tensors holding - // the previous chunks' pre-projected attention keys/values (bias - // included). When set, the block computes K/V projections only for - // the new frames, concatenates the cache in front, runs attention - // against the combined window, and emits the rotated cache via - // ggml_cpy into the *_out tensors (the driver copies them into the - // persistent cache after compute, exactly like streaming_channel_*). - // streaming_channel_in/_out are ignored in this mode — the channel - // cache is the recompute-K/V representation of the same state. + // Optional streaming KV cache (all four set together, or none): + // persistent [d_model, T_cache] pre-projected keys/values. When set, K/V + // are projected only for the new frames, the cache is prepended, and the + // rotated cache is emitted into the *_out tensors. streaming_channel_in/ + // _out are ignored in this mode (same state, recompute-K/V form). ggml_tensor * streaming_kv_k_in = nullptr; ggml_tensor * streaming_kv_v_in = nullptr; ggml_tensor * streaming_kv_k_out = nullptr; ggml_tensor * streaming_kv_v_out = nullptr; - // Optional precomputed rel-pos projection for this block: the - // [head_dim, pos_len, n_head, 1] result of - // cont(permute(reshape(attn_pos_w @ pos_emb))) from a previous - // chunk with identical geometry. When non-null, rel_pos_mhsa - // consumes it directly and never touches pos_emb (which may then - // be null). Streaming-only memoization; offline paths leave this - // null. + // Optional precomputed rel-pos projection [head_dim, pos_len, n_head, 1] + // from a prior chunk of identical geometry; when non-null rel_pos_mhsa + // consumes it directly and ignores pos_emb. Streaming-only memoization. ggml_tensor * streaming_pos_proj_in = nullptr; - // Graph the helpers use to forward-expand their cache-write cpy - // nodes. The streaming cache outputs are SIDE outputs (not - // reachable from the encoder's `out` tensor), so they have to be - // added to the graph explicitly. The caller sets this to the - // graph it'll later run on; the helpers do - // ggml_build_forward_expand(streaming_graph, cpy_node) when they - // emit a cache write. Required when streaming_*_out is non-null. + // Graph the helpers forward-expand their cache-write cpy nodes onto. The + // streaming cache outputs are SIDE outputs (not reachable from the + // encoder's `out`), so they must be added explicitly. Required when + // streaming_*_out is non-null. ggml_cgraph * streaming_graph = nullptr; }; -// =========================================================================== -// Low-level helpers (exposed for family-specific hand-built blocks) -// =========================================================================== +// Low-level helpers (exposed for family-specific hand-built blocks). // Attach a debug name to a tensor. Safe on nullptr. ggml_tensor * named(ggml_tensor * t, const char * name); @@ -409,9 +281,8 @@ ggml_tensor * macaron_ff_residual(ggml_context * ctx, // matrix rotated so column k holds the score for relative offset k. ggml_tensor * rel_shift(ggml_context * ctx, ggml_tensor * x); -// f32-friendly Conv1D. Use this instead of ggml_conv_1d for fp32 -// kernels — see the long comment above conv_2d_dw_f32 in -// conformer.cpp for the Metal backstory. +// f32-friendly Conv1D. Use instead of ggml_conv_1d for fp32 kernels on +// Metal (see conv_2d_dw_f32 in conformer.cpp). ggml_tensor * conv_1d_f32(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data, @@ -419,10 +290,8 @@ ggml_tensor * conv_1d_f32(ggml_context * ctx, int padding, int dilation); -// f32-friendly 2D depthwise conv. Use this instead of -// ggml_conv_2d_dw / ggml_conv_2d_dw_direct for fp32 kernels when the -// compute may run on Metal — the full story is in the comment above -// the implementation. +// f32-friendly 2D depthwise conv. Use instead of ggml_conv_2d_dw / +// ggml_conv_2d_dw_direct for fp32 kernels on Metal (see conformer.cpp). ggml_tensor * conv_2d_dw_f32(ggml_context * ctx, ggml_tensor * kernel, ggml_tensor * data, @@ -456,9 +325,7 @@ ggml_tensor * add_conv_bias(ggml_context * ctx, ggml_tensor * conv_out, ggml_tensor * bias_1d); -// =========================================================================== -// Sub-blocks (exposed for families that hand-build a block) -// =========================================================================== +// Sub-blocks (exposed for families that hand-build a block). // Convolution sub-block: pointwise -> GLU -> optional valid-frame mask // -> depthwise -> BN/LN -> SiLU -> pointwise. Operates on post-LayerNorm @@ -510,31 +377,14 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, ggml_tensor * k_full = nullptr, ggml_tensor * v_full = nullptr); -// =========================================================================== -// Top-level block + pre-encode -// =========================================================================== +// Top-level block + pre-encode. -// Optional per-block observer hook. `build_conformer_block` calls -// `on_point` at the five canonical residual / post-LN points during -// the block forward, with `tag` set to a compile-time constant -// string: -// -// "after_ff1" — after the FF1 macaron residual merge -// "after_attn" — after the attention residual merge -// "after_conv" — after the conv module residual merge -// "after_ff2" — after the FF2 macaron residual merge -// "out" — after the final per-block LayerNorm (== block out) -// -// Callers typically use this to attach a ggml name, stash the -// tensor in a dump-slot struct, and/or mark the tensor for graph -// output. The observer is the supported replacement for -// hand-building a block out of the exposed sub-helpers just to get -// named dump points: it keeps the block implementation single-source -// and removes the lockstep maintenance hazard. -// -// Pointer-of-struct + user pointer (not std::function) keeps this -// header C-ABI-friendly and free of . Pass nullptr for -// the observer argument to skip observation entirely. +// Optional per-block observer hook. `build_conformer_block` calls `on_point` +// at five canonical residual / post-LN points (compile-time `tag`): +// "after_ff1", "after_attn", "after_conv", "after_ff2", "out". +// Callers attach a ggml name, stash the tensor, and/or mark it for graph +// output. Pointer-of-struct + user pointer (not std::function) keeps the +// header C-ABI-friendly. Pass nullptr to skip observation. struct BlockObserver { void (*on_point)(void * user, const char * tag, @@ -555,22 +405,15 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, const BlockObserver * obs = nullptr); // Per-stage valid-frame masks for variable-length offline batching of the -// pre_encode conv stem. Zero-padding utterances to a common length is NOT -// transparent through this conv stack: the convs carry bias + ReLU, so the -// padded region becomes non-zero after the first conv and then leaks into -// each utterance's last VALID frame via the next stride-2 conv's receptive -// field — corrupting exactly the boundary frame (it flips trailing CTC -// tokens). The fix (NeMo's masked subsampling) is to zero each conv -// intermediate's padded time region after every ReLU so the next conv sees -// zeros beyond each utterance's boundary, exactly like a standalone run. -// -// When a non-null PreEncodeValidMasks is passed, build_pre_encode allocates -// three time-valid mask graph inputs (one per ReLU stage) sized -// ne=[1, H_stage, 1, B] (broadcast over freq + channels), applies them, and -// writes the handles back here for the driver to fill (1.0 on valid frames, -// 0.0 on padded). Only meaningful for the non-causal (offline) pre_encode; -// the caller should skip it for causal/streaming variants whose output-length -// formula differs. +// pre_encode conv stem. Zero-padding to a common length is NOT transparent +// through this stack: convs carry bias + ReLU, so the padded region becomes +// non-zero after the first conv and then leaks into each utterance's last +// VALID frame via the next stride-2 conv's receptive field (flips trailing +// CTC tokens). NeMo's masked subsampling zeros each conv intermediate's +// padded time region after every ReLU. When PreEncodeValidMasks is non-null, +// build_pre_encode allocates three mask graph inputs ne=[1, H_stage, 1, B] +// (one per ReLU stage), applies them, and writes the handles back for the +// driver to fill. Offline (non-causal) pre_encode only. struct PreEncodeValidMasks { ggml_tensor * mask_s1 = nullptr; // after relu0 ggml_tensor * mask_s2 = nullptr; // after relu3 @@ -579,15 +422,10 @@ struct PreEncodeValidMasks { // DwStridingSubsampling pre_encode stack. Returns the final // [d_model, T_enc, 1, 1] tensor or nullptr on a shape-sanity failure. -// `name_prefix`, if non-null, is used to attach ggml_set_name() calls -// to every intermediate (pattern: ".conv0", ".relu0", -// ...). Pass nullptr to skip naming entirely. -// -// Uses pe.out_w->ne[0] as a final d_model sanity check. `error_tag` -// supplies the family name used in the diagnostic on shape mismatch. -// -// valid_masks (optional): when non-null, enables the masked-subsampling path -// above and is filled with the three mask input handles. +// `name_prefix` (if non-null) names every intermediate ".conv0" +// etc. Uses pe.out_w->ne[0] as a d_model sanity check; `error_tag` names the +// family in the diagnostic. valid_masks (optional) enables the +// masked-subsampling path above and is filled with the mask input handles. ggml_tensor * build_pre_encode(ggml_context * ctx, const PreEncodeView & pe, ggml_tensor * mel_in, diff --git a/src/granite_conformer/shaw_attn.cpp b/src/granite_conformer/shaw_attn.cpp index 9096ff04..6a047623 100644 --- a/src/granite_conformer/shaw_attn.cpp +++ b/src/granite_conformer/shaw_attn.cpp @@ -1,7 +1,6 @@ // src/granite_conformer/shaw_attn.cpp - shared Shaw block-local attention. // -// See shaw_attn.h for the contract. Both Granite families' encoders -// inline this op identically; this file is the single source of truth. +// See shaw_attn.h for the contract. #include "shaw_attn.h" #include "transcribe-log.h" @@ -42,13 +41,10 @@ ggml_tensor * shaw_block_attn(ggml_context * ctx, const int64_t T_pad = static_cast(context_size) * num_blocks; const float scale = 1.0f / std::sqrt(static_cast(head_dim)); - // Offline batch: x ne=[d_model, T_enc, B]. The block-local attention - // treats each (time-block) independently, so B utterances fold into the - // block axis: num_blocks_eff = num_blocks * B. Memory order is - // [.., T_pad(=ctx*num_blocks), B], so a single reshape splits the time - // axis into (ctx, num_blocks) with B outermost == (ctx, num_blocks*B). - // B == 1 (AR granite, single-shot) keeps num_blocks_eff == num_blocks and - // is byte-identical to the pre-batch graph. The caller tiles the per-block + // Offline batch: x ne=[d_model, T_enc, B]. Block-local attention treats + // each time-block independently, so B utterances fold into the block + // axis: num_blocks_eff = num_blocks * B (memory order [.., T_pad, B] + // splits cleanly into (ctx, num_blocks*B)). The caller tiles the per-block // pad mask across B and sizes zero_pad with the batch. const int64_t B = x->ne[2]; const int64_t num_blocks_eff = static_cast(num_blocks) * B; @@ -101,11 +97,10 @@ ggml_tensor * shaw_block_attn(ggml_context * ctx, k = reshape_qkv(k); v = reshape_qkv(v); - // ----- Shaw positional bias ----- - // rel_pos_emb is [head_dim, 2*max_pos_emb+1]. attention_dists is - // [context_size, context_size] int32. Materialise the per-(c, r) - // lookup as a [head_dim, context_size*context_size] tensor, then - // reshape to [head_dim, r=context_size, c=context_size]. + // Shaw positional bias. rel_pos_emb [head_dim, 2*max_pos_emb+1], dists + // [context_size, context_size] int32. Materialise the per-(c, r) lookup + // as [head_dim, context_size*context_size], then reshape to + // [head_dim, r=context_size, c=context_size]. ggml_tensor * dists_flat = ggml_reshape_1d( ctx, dists, static_cast(context_size) * context_size); ggml_tensor * rel_lookup = ggml_get_rows(ctx, w.attn_rel_pos_emb, @@ -119,10 +114,8 @@ ggml_tensor * shaw_block_attn(ggml_context * ctx, ggml_tensor * q_perm = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); ggml_tensor * pos_attn = ggml_mul_mat(ctx, rel_lookup, q_perm); - // ----- QK^T ----- - // ggml_mul_mat(K, Q) with K, Q both [head_dim, context_size, - // n_heads, num_blocks] -> [context_size=k, context_size=q, - // n_heads, num_blocks]. + // QK^T: ggml_mul_mat(K, Q), both [head_dim, context_size, n_heads, + // num_blocks] -> [k=context_size, q=context_size, n_heads, num_blocks]. ggml_tensor * kq = ggml_mul_mat(ctx, k, q); // Align pos_attn's axis order to kq's before the add. @@ -147,8 +140,7 @@ ggml_tensor * shaw_block_attn(ggml_context * ctx, ggml_tensor * v_t = ggml_cont(ctx, ggml_permute(ctx, v, 1, 0, 2, 3)); ggml_tensor * out = ggml_mul_mat(ctx, v_t, attn); - // Reshape back to [inner_dim, T_pad, B] (B==1 collapses to the - // pre-batch [inner_dim, T_pad]). + // Reshape back to [inner_dim, T_pad, B]. out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); out = ggml_reshape_3d(ctx, out, inner_dim, T_pad, B); diff --git a/src/granite_conformer/shaw_attn.h b/src/granite_conformer/shaw_attn.h index a695c619..1d21e7da 100644 --- a/src/granite_conformer/shaw_attn.h +++ b/src/granite_conformer/shaw_attn.h @@ -1,23 +1,19 @@ // src/granite_conformer/shaw_attn.h - shared Shaw block-local attention. // -// Reused by both `src/arch/granite/encoder.cpp` (AR audio-LLM) and -// `src/arch/granite_nar/encoder.cpp` (NAR editor). The two families -// share an audio encoder up to the projector; this is the only -// Granite-specific Conformer op without an analog in `src/conformer/`. +// Used by src/arch/granite/encoder.cpp (AR audio-LLM) and +// src/arch/granite_nar/encoder.cpp (NAR editor); the only Granite-specific +// Conformer op without an analog in src/conformer/. // -// Mirrors `GraniteSpeechConformerAttention.forward`: +// Mirrors GraniteSpeechConformerAttention.forward: // // h = pre_norm(x) # at T_enc // if remainder > 0: h = F.pad(h, (0, 0, 0, pad)) # zero-pad to T_pad // q = to_q(h); kv = to_kv(h); k, v = kv.chunk(2, dim=-1) // ... block-local attention with Shaw bias + last-block pad mask ... -// out = out.reshape(..., T_pad, inner_dim) // out = to_out(out[:, :num_features, :]) # slice back to T_enc -// return out # at T_enc // -// The slice happens BEFORE to_out (upstream order), so the following -// conv module's depthwise kernel never pulls pad-row garbage into -// valid frames near the tail. +// The slice happens BEFORE to_out (upstream order), so the following conv +// module's depthwise kernel never pulls pad-row garbage into valid frames. #pragma once @@ -25,10 +21,9 @@ namespace transcribe::granite_conformer { -// Per-block Shaw attention weights. Both `GraniteEncBlock` and -// `GraniteNarEncBlock` carry the same set of tensors with the same -// field names; this packs them so the helper has no compile-time -// dependency on either family's weights struct. +// Per-block Shaw attention weights. Field names match GraniteEncBlock and +// GraniteNarEncBlock, packed so the helper has no compile-time dependency on +// either family's weights struct. struct ShawAttnWeights { ggml_tensor * norm_attn_w = nullptr; // [hidden] ggml_tensor * norm_attn_b = nullptr; // [hidden] (nullable) @@ -41,23 +36,29 @@ struct ShawAttnWeights { // Build the Shaw block-local attention sub-graph. // +// Batch-capable: the input carries an optional utterance batch B on ne[2] +// (B==1 for single-shot). Block-local attention treats each time-block +// independently, so the B utterances fold into the block axis internally +// (num_blocks_eff = num_blocks * B); the caller tiles the per-block pad +// mask across B and sizes zero_pad with the batch. +// // Inputs: -// x : [d_model, T_enc] pre-norm input. -// zero_pad : [d_model, T_pad - T_enc] f32 zeros, or nullptr when +// x : [d_model, T_enc, B] pre-norm input (B==1 single-shot). +// zero_pad : [d_model, T_pad - T_enc, B] f32 zeros, or nullptr when // T_enc % context_size == 0. // dists : [context_size * context_size] int32 Shaw lookup -// indices (`clamp(k - q + max_pos_emb, 0, 2*max_pos_emb)`). -// pad_mask_3d : [context_size, context_size, num_blocks] f32 additive -// mask (all-zero except the last slice, which has -INF -// on pad-K columns). +// indices (`clamp(c - r + max_pos_emb, 0, 2*max_pos_emb)`, +// c = key/column, r = query/row). Shared across B. +// pad_mask_3d : [context_size, context_size, num_blocks * B] f32 additive +// mask (all-zero except each utterance's last block slice, +// which has -INF on pad-K columns). The caller tiles the +// per-block pattern across B. // w : per-block weights (see ShawAttnWeights). -// n_heads, head_dim, context_size, num_blocks, T_enc : shape scalars. -// layer_norm_eps : eps for the pre-norm (both Granite families use -// 1e-5; expose as a parameter so callers stay in -// control). +// n_heads, head_dim, context_size, num_blocks, T_enc : shape scalars +// (num_blocks is per-utterance; B is read from x->ne[2]). +// layer_norm_eps : pre-norm eps (both Granite families use 1e-5). // -// Returns: [d_model, T_enc] (matches input shape). Returns nullptr on -// shape error (with a stderr diagnostic). +// Returns: [d_model, T_enc, B] (matches input), or nullptr on shape error. ggml_tensor * shaw_block_attn(ggml_context * ctx, ggml_tensor * x, ggml_tensor * zero_pad, diff --git a/src/sanm/sanm.cpp b/src/sanm/sanm.cpp index fa557686..e8e9ac37 100644 --- a/src/sanm/sanm.cpp +++ b/src/sanm/sanm.cpp @@ -12,17 +12,10 @@ namespace transcribe::sanm { -// SAN-M scales activations by sqrt(d_model) before the encoder stack, so its -// internal activations run large. With F16 weights, CUDA's cuBLAS path -// accumulates matmuls in F16 (CUBLAS_COMPUTE_16F) and those large dot-product -// partial sums overflow F16's ~65504 range -> NaNs / garbage. Metal and CPU -// accumulate in F32 regardless, which is why F16 only breaks on CUDA. Forcing -// GGML_PREC_F32 makes cuBLAS accumulate in F32 (matching the CPU reference). -// -// Gated on F16 weights so it cannot perturb any already-validated numerics: -// CPU/Metal ignore the prec hint (always F32-accumulate), and BF16/quantized -// weights are left on their default path (BF16 already COMPUTE_32F; quantized -// goes through MMQ). Only CUDA-F16 changes. +// Force F32 accumulation for F16 weights (see mul_mat_f32acc in +// causal_lm.cpp for the CUDA COMPUTE_16F saturation rationale). SAN-M scales +// activations by sqrt(d_model), so its internal activations run large and +// overflow F16 on CUDA without this. static ggml_tensor * mul_mat_f32acc(ggml_context * ctx, ggml_tensor * w, ggml_tensor * x) @@ -90,7 +83,7 @@ ggml_tensor * fsmn_branch(ggml_context * ctx, // [T, 1, d_model, B] -> [T, d_model, B]. fsmn = ggml_reshape_3d(ctx, fsmn, fsmn->ne[0], d_model, B); } else { - // Single-shot: keep the validated im2col path bit-for-bit. + // Single-shot: im2col path. fsmn = conf::conv_1d_dw_f32(ctx, fsmn_w, v_t, /*stride=*/1, /*padding=*/padding, /*dilation=*/1); @@ -142,12 +135,11 @@ ggml_tensor * sanm_attention(ggml_context * ctx, // a strided alias of V. ggml_tensor * v_pre = ggml_cont(ctx, v); - // ----- FSMN branch (parallel to SDPA) ----- + // FSMN branch (parallel to SDPA). ggml_tensor * fsmn = fsmn_branch(ctx, v_pre, b.attn_fsmn_w, kernel, p.conv_pad_mask); - // ----- SDPA ----- - // Reshape Q,K,V to [head_dim, n_heads, T, B]. The fused QKV split above + // 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. auto split_heads = [&](ggml_tensor * t) { t = ggml_cont(ctx, t); @@ -167,12 +159,10 @@ ggml_tensor * sanm_attention(ggml_context * ctx, kh = to_attn_layout(kh); vh = to_attn_layout(vh); - // A variable-length batch supplies an additive key-padding mask; the - // flash kernel's masked batched form is not exercised here, so route + // A variable-length batch supplies an additive key-padding mask; route // through the manual mul_mat + soft_max path, which broadcasts the - // [T_k, 1, 1, B] mask cleanly over queries and heads. Same-length - // batches (and single-shot) take the flash path, bit-identical to the - // pre-batch graph at B == 1. + // [T_k, 1, 1, B] mask cleanly over queries and heads. Same-length batches + // (and single-shot) take the flash path. ggml_tensor * o; if (p.use_flash && p.attn_pad_mask == nullptr) { o = ggml_flash_attn_ext( @@ -197,7 +187,6 @@ ggml_tensor * sanm_attention(ggml_context * ctx, o = ggml_reshape_3d(ctx, o, d_model, T, B); } - // Output projection. o = mul_mat_f32acc(ctx, b.attn_out_w, o); o = ggml_add(ctx, o, b.attn_out_b); diff --git a/src/sanm/sanm.h b/src/sanm/sanm.h index 529ee580..f9791104 100644 --- a/src/sanm/sanm.h +++ b/src/sanm/sanm.h @@ -1,34 +1,14 @@ // src/sanm/sanm.h - shared SAN-M encoder helpers + sinusoidal PE. // -// Both SenseVoice (`SenseVoiceEncoderSmall`) and Fun-ASR-Nano -// (`FunASRNano.encoder`) ship identical SAN-M block code: a fused QKV -// projection with a parallel FSMN depthwise-conv1d branch on V, SDPA -// over the QKV split, and a 2-layer ReLU FFN. The encoder structures -// only differ in the surrounding shape — sensevoice prepends prefix -// embeddings and adds a CTC head, funasr_nano feeds the encoder output -// directly into an audio adaptor. +// SenseVoice (SenseVoiceEncoderSmall) and Fun-ASR-Nano (FunASRNano.encoder) +// ship identical SAN-M block code: a fused QKV projection with a parallel +// FSMN depthwise-conv1d branch on V, SDPA over the QKV split, and a 2-layer +// ReLU FFN. Per-block primitives follow the conformer-style view + +// free-function pattern (src/conformer/conformer.h); sub-blocks are exposed +// for families that instrument dump points. build_sinusoidal_pe is the +// host-side PE-table companion both families add before the first block. // -// This module exposes the per-block primitives via the conformer-style -// view + free-function pattern (`src/conformer/conformer.h`): -// -// - SanmBlockView : nullable-pointer projection over one block's -// weights. Each family fills it from its own -// per-block weight struct at the call site. -// - SanmBlockParams : per-call shape (n_heads, d_model, kernel). -// - sanm_block_residual / sanm_block_projection : the two block -// variants used by SenseVoice / Fun-ASR-Nano. -// - sub-blocks (sv_layer_norm, fsmn_branch, sanm_attention, -// sanm_ffn) are exposed too for families that want to instrument -// intermediate dump points. -// -// `build_sinusoidal_pe` is the host-side companion that builds the -// 1-based-positions PE table. Both families add the table to the -// frontend output before the first SAN-M block; the consumer set is -// identical so the function lives here. Promote to a standalone -// `transcribe-sinusoidal-pe.{h,cpp}` only if a third consumer needs PE -// without SAN-M. -// -// LayerNorm epsilon is 1e-12 (FunASR's `transformer.layer_norm.LayerNorm`), +// LayerNorm epsilon is 1e-12 (FunASR's transformer.layer_norm.LayerNorm), // NOT the conformer-shared 1e-5. #pragma once @@ -40,14 +20,12 @@ struct ggml_tensor; namespace transcribe::sanm { -// LayerNorm epsilon used by FunASR's `transformer.layer_norm.LayerNorm`. -// Hardcoded here to match the reference; if a future SAN-M variant -// changes it, plumb through SanmBlockParams rather than bumping this. +// LayerNorm epsilon (FunASR's transformer.layer_norm.LayerNorm); hardcoded +// to match the reference. constexpr float kLayerNormEps = 1e-12f; -// Nullable-pointer projection over one SAN-M block's weights. Field -// names match SenseVoice's `SenseVoiceBlock` and Fun-ASR-Nano's -// `EncBlock` so a family can build the view with a single helper. +// Nullable-pointer projection over one SAN-M block's weights. Field names +// match SenseVoice's SenseVoiceBlock and Fun-ASR-Nano's EncBlock. // // Tensor shapes (ggml channel-innermost layout): // norm_attn_w/b [d_in] @@ -77,22 +55,15 @@ struct SanmBlockView { ggml_tensor * ffn_fc2_b = nullptr; }; -// Per-call shape for the SAN-M block. -// -// The batch axis rides the activation's ne[2]: a single-utterance call -// passes x ne=[d, T] (ne[2] == 1) and every helper below derives B from -// x->ne[2], so the B == 1 path is byte-identical to the pre-batch code. -// Offline batching (transcribe_run_batch) packs B utterances at ne[2] and, -// when their lengths differ, hands in the two padding masks so a padded -// tail can never perturb a real frame: +// Per-call shape for the SAN-M block. The batch axis rides ne[2] (B derived +// from x->ne[2]). Offline batching with mismatched lengths supplies two +// padding masks so a padded tail never perturbs a real frame: // - attn_pad_mask : additive key-padding mask ne=[T_k, 1, 1, B] f32 -// (0 on real keys, -INF on padded). Added onto the attention score -// matrix. Its presence also forces the manual SDPA path (the flash -// kernel's masked batched form is not exercised here). -// - conv_pad_mask : FSMN valid-frame mask ne=[1, T, B] f32 (1 on real -// frames, 0 on padded). Multiplies V before the depthwise conv so a -// padded frame contributes 0 to its real neighbours. -// Same-length batches leave both null and run the bit-exact flash path. +// (0 on real keys, -INF on padded), added onto the score matrix; its +// presence forces the manual SDPA path. +// - conv_pad_mask : FSMN valid-frame mask ne=[1, T, B] f32 (1 real, 0 pad), +// multiplies V before the depthwise conv. +// Same-length batches leave both null and run the flash path. struct SanmBlockParams { int n_heads = 0; int d_model = 0; diff --git a/src/transcribe-abi.h b/src/transcribe-abi.h index 661eb4f3..bdf46770 100644 --- a/src/transcribe-abi.h +++ b/src/transcribe-abi.h @@ -25,11 +25,10 @@ inline transcribe_status check_struct_size(uint64_t got, uint64_t want) { return TRANSCRIBE_OK; } -// Kept as an alias for call-site readability ("this is an input params -// struct"). Behavior is identical to check_struct_size: in pre-1.0 the -// `0 means defaults` relaxation has been removed so that uninitialized -// stack memory with a coincidentally-zero struct_size byte stops being -// silently accepted as a defaults shortcut. Defaults come from NULL. +// Alias for call-site readability ("this is an input params struct"). +// Behavior is identical to check_struct_size: struct_size == 0 is rejected +// (so coincidentally-zero stack memory is not taken as a defaults +// shortcut); defaults come from NULL. inline transcribe_status check_input_struct_size(uint64_t got, uint64_t want) { return check_struct_size(got, want); } diff --git a/src/transcribe-arch.cpp b/src/transcribe-arch.cpp index fbbfdaa0..84e72daf 100644 --- a/src/transcribe-arch.cpp +++ b/src/transcribe-arch.cpp @@ -1,15 +1,8 @@ // transcribe-arch.cpp - explicit registry of supported architectures. // -// The registry is intentionally a function-local static array. Pros: -// - No static-initializer-order dependency between this TU and the -// per-family TUs that define each Arch instance (the array holds -// pointers, not copies, and is initialized lazily on first call). -// - The full set of supported architectures is grep-able from one file. -// - Adding a family is a two-line edit (forward declaration + array -// entry) plus the per-family TU. -// -// Cons (acceptable for now): -// - Linear search. Single-digit n; lookups happen once per model load. +// The registry is a function-local static array of pointers (no static- +// initializer-order dependency on the per-family TUs, grep-able from one +// file). Adding a family is a two-line edit plus the per-family TU. #include "transcribe-arch.h" diff --git a/src/transcribe-arch.h b/src/transcribe-arch.h index d3c1ad80..99781996 100644 --- a/src/transcribe-arch.h +++ b/src/transcribe-arch.h @@ -1,24 +1,11 @@ // transcribe-arch.h - per-family architecture trait + registry. // -// This header is INTERNAL. The public C ABI in include/transcribe.h knows -// nothing about Arch; it dispatches through find_arch() inside -// transcribe_model_load_file and the rest of the lifecycle entry points. -// -// Each model family (parakeet, eventually whisper, sense-voice, etc.) -// provides exactly one transcribe::Arch instance whose function pointers -// implement that family's load / init_context / run behavior. The -// registry is an explicit array (no static initializers, no dlopen, no -// global registration tricks) so the set of supported architectures is -// grep-able from one file. -// -// Adding a new family is a two-line edit in transcribe-arch.cpp plus a -// new src/arch//model.cpp that defines the Arch instance. -// -// The Arch trait contains only the things that genuinely vary per -// family. Lifecycle teardown (free_model / free_context) is handled -// by virtual destructors on transcribe_model / transcribe_session, -// and the introspection accessors (capabilities, backend_string, -// arch_string, variant_string) read directly from the base struct. +// INTERNAL. The public C ABI knows nothing about Arch; it dispatches +// through find_arch(). Each family provides exactly one transcribe::Arch +// instance whose function pointers implement its load / init_context / run +// behavior; the registry is an explicit array in transcribe-arch.cpp. The +// trait holds only what genuinely varies per family: teardown is handled by +// the base virtual destructors and introspection reads the base struct. #pragma once @@ -28,32 +15,18 @@ namespace transcribe { class Loader; -// Per-family trait. Function pointers may be null if the corresponding -// entry point is not yet implemented for that family; the central -// dispatch in transcribe.cpp converts null entries into -// TRANSCRIBE_ERR_NOT_IMPLEMENTED so a partially-wired family does not -// crash mid-run. +// Per-family trait. Function pointers may be null when an entry point is +// not yet implemented; the central dispatch converts null entries into +// TRANSCRIBE_ERR_NOT_IMPLEMENTED. // // Ownership conventions: -// -// load: the Loader has already opened the GGUF and read the -// architecture/variant strings. The handler may call -// loader.release_gguf() to take ownership of the gguf_context, in which -// case it is responsible for freeing it (typically in the per-family -// model's destructor). If the handler does not call release_gguf(), -// the Loader's destructor frees the context on stack unwinding. -// -// On success, *out_model points at a heap-allocated derived -// transcribe_model. On failure, *out_model is left null and the -// handler returns a non-OK status. -// -// init_context: same convention. On success *out_ctx points at a -// heap-allocated derived transcribe_session. The dispatcher has -// already validated that model and params are non-null. -// -// Cleanup is NOT in the trait: the central dispatcher's -// transcribe_model_free / transcribe_session_free use `delete` against -// the base, and the virtual destructors do the right thing. +// load / init_context: on success *out_model / *out_ctx points at a +// heap-allocated derived object; on failure it is left null and a +// non-OK status returned. load may call loader.release_gguf() to take +// ownership of the gguf_context (and then free it, typically in the +// model destructor); otherwise the Loader frees it on unwinding. +// Cleanup is NOT in the trait: transcribe_model_free / +// transcribe_session_free `delete` against the base virtual destructor. struct Arch { // The string compared against general.architecture in the GGUF KV. // Must be a NUL-terminated string with static lifetime. @@ -76,39 +49,24 @@ struct Arch { const transcribe_run_params * params); // Optional offline batch fast path. When non-null, transcribe_run_batch - // delegates here to process all `n` utterances in a single dispatch - // (e.g. a batched encoder graph) for device-level throughput. When - // null, the central dispatcher falls back to calling run() once per - // utterance and snapshotting each result, so EVERY family supports the - // batch API for correctness regardless of whether it wires this hook; - // a family opts into the fast path by implementing it. + // delegates here to process all `n` utterances in one dispatch; when + // null the dispatcher falls back to calling run() once per utterance, + // so every family supports the batch API regardless. // // Contract: - // - The dispatcher has already validated the shared run_params - // (struct size, _RUN-slot ext shape/kind, enum range, timestamp - // ceiling, language, TRANSLATE support, run_validate) ONCE and - // cleared both the scratch result slot and ctx->batch_results. - // - pcm[i] / n_samples[i] are the i-th utterance (i in [0, n)). - // 16 kHz mono float32, same as run(). The hook is responsible for - // per-utterance input validation (pcm[i] == null or - // n_samples[i] <= 0 is a per-utterance failure, recorded as a - // non-OK ResultSet::status with has_result == false, NOT a - // whole-batch error). - // - On an OK return the hook MUST push exactly `n` entries onto - // ctx->batch_results, one per utterance in order, and SHOULD leave - // the scratch slot mirroring batch_results[0] on return (the - // dispatcher also enforces this). Poll ctx->poll_abort() between - // utterances / decode steps; on abort, push the utterances retained - // as completed results and return TRANSCRIBE_ERR_ABORTED — the hook - // need NOT pad missing slots. The dispatcher pads batch_results back - // up to `n` with TRANSCRIBE_ERR_ABORTED ResultSets (status meaning - // "did not complete because the batch was aborted"), so after an OK - // or ABORTED return batch_results.size() == n. - // - The function return value is the whole-batch status: OK when the - // dispatch ran (per-utterance failures live in each ResultSet), - // non-OK only for whole-batch faults (OOM, abort, backend error). - // For a whole-batch fault other than abort (OOM, backend error) the - // result count is unspecified; the non-OK return is the signal. + // - The dispatcher has validated the shared run_params ONCE and + // cleared the scratch slot and ctx->batch_results. + // - pcm[i] / n_samples[i] are the i-th utterance (16 kHz mono f32). + // The hook does per-utterance input validation (a bad utterance is + // a non-OK ResultSet with has_result == false, not a batch error). + // - On OK the hook MUST push exactly `n` entries onto batch_results + // in order and SHOULD leave the scratch slot mirroring + // batch_results[0]. Poll ctx->poll_abort() between utterances; on + // abort, push the completed results and return TRANSCRIBE_ERR_ABORTED + // (the dispatcher pads batch_results back up to `n`). + // - The return value is the whole-batch status: OK when the dispatch + // ran (per-utterance failures live in each ResultSet), non-OK only + // for whole-batch faults (OOM, abort, backend error). transcribe_status (*run_batch)( struct transcribe_session * ctx, const float * const * pcm, @@ -117,58 +75,27 @@ struct Arch { const transcribe_run_params * params); // Optional streaming hooks. stream_begin / stream_feed / - // stream_finalize form a required triple: a family that wants to - // support streaming MUST wire all three. The dispatcher checks - // the triple at begin time and returns NOT_IMPLEMENTED if any of - // them is NULL (a partially-wired family must never let a caller - // enter ACTIVE and then get stuck on the next call). stream_reset - // is optional — when NULL, transcribe_stream_reset still wipes - // dispatcher-owned state and runs the clear_result path, which - // is sufficient when a family doesn't need to release per- - // utterance buffers explicitly. + // stream_finalize are a required triple (the dispatcher returns + // NOT_IMPLEMENTED if any is NULL); stream_reset is optional (when NULL + // the dispatcher still wipes its own state and runs clear_result). // - // Hook responsibilities: - // - // stream_validate (optional): pre-flight validation of family- - // specific parameters and stream_params->family extension values. - // Called BEFORE the dispatcher clears the previous result - // snapshot and BEFORE the lifecycle moves out of its current - // state. Must be pure: zero mutation of ctx or model. On non-OK - // return the dispatcher returns the status to the caller with - // the previous result snapshot fully preserved and the stream - // lifecycle untouched (no FAILED transition). Families that - // accept extension structs SHOULD validate enum-from-menu and - // value-shape rules here so a caller typo does not destroy the - // prior utterance's transcript. - // - // stream_begin: install per-utterance state on the derived - // context. The central dispatcher has already run - // stream_validate (if any), cleared the result snapshot, and - // set ctx->stream_state to ACTIVE. On non-OK return the - // dispatcher transitions to FAILED and preserves the status in - // stream_last_status. Families MAY repeat cheap validation - // here as defense in depth, but the canonical place to reject - // bad caller input is stream_validate. - // - // stream_feed: consume the PCM (n_samples may be 0) and update - // the result vectors / committed counts / audio cursors on - // ctx in place. When update is non-null, fill it with the - // family's per-call change metadata; the dispatcher has - // already zero-initialized it. Poll ctx->poll_abort() at - // chunk and decode-step boundaries; if it fires, return - // TRANSCRIBE_ERR_ABORTED (the dispatcher will transition to - // FAILED and record the status, and transcribe_was_aborted - // will distinguish abort from other failures). - // - // stream_finalize: flush buffered audio, satisfy lookahead, - // emit remaining text, mark all output as committed. The - // dispatcher transitions to FINISHED on TRANSCRIBE_OK and - // FAILED otherwise. - // - // stream_reset: release the family's per-utterance state and - // buffered audio contents (keeping allocated buffers for - // reuse). The dispatcher clears the result snapshot and - // forces stream_state back to IDLE around this call. + // stream_validate (optional): pure pre-flight validation of family + // params / stream_params extension, called BEFORE the result + // snapshot is cleared and the lifecycle moves. On non-OK the prior + // snapshot and lifecycle are fully preserved (no FAILED). The right + // place to vet caller input so a typo can't destroy a transcript. + // stream_begin: install per-utterance state. The dispatcher has + // already run stream_validate, cleared the snapshot, and set state + // to ACTIVE; on non-OK it transitions to FAILED and records the + // status in stream_last_status. + // stream_feed: consume the PCM (n_samples may be 0), update result + // vectors / committed counts / cursors in place, fill `update` when + // non-null. Poll ctx->poll_abort(); on abort return + // TRANSCRIBE_ERR_ABORTED. + // stream_finalize: flush buffered audio, emit remaining text, mark + // all output committed. Dispatcher -> FINISHED on OK, FAILED else. + // stream_reset: release per-utterance state / buffered audio (keeping + // allocations); the dispatcher forces state back to IDLE around it. transcribe_status (*stream_validate)( const struct transcribe_session * ctx, const transcribe_run_params * run_params, @@ -192,62 +119,35 @@ struct Arch { void (*stream_reset)( struct transcribe_session * ctx); - // Optional kind+slot probe. Returns true when the loaded model - // variant accepts the named extension kind on the given slot - // (transcribe_run_params::family for _RUN, transcribe_stream_params::family - // for _STREAM). NULL means "no family extension kinds accepted in - // any slot" — the dispatcher's transcribe_model_accepts_ext_kind() - // returns false in that case. - // - // Acceptance is per-loaded-variant AND per-slot: - // - parakeet returns true for (_STREAM, PARAKEET_STREAM) on - // cache-aware variants and for (_STREAM, PARAKEET_BUFFERED_STREAM) - // on chunked-attention variants — never both on the same loaded - // model — and always false for the _RUN slot. - // - moonshine_streaming returns true for - // (_STREAM, MOONSHINE_STREAMING_STREAM) only. - // - whisper returns true for (_RUN, WHISPER_RUN) only. - // A family that genuinely has no extension surface for a slot must - // return false rather than NULL'ing this hook; NULL is reserved for - // "no extension surface at all." + // Optional kind+slot probe. Returns true when the loaded model variant + // accepts the named extension kind on the given slot (run_params::family + // for _RUN, stream_params::family for _STREAM). Acceptance is + // per-loaded-variant AND per-slot. NULL is reserved for "no extension + // surface at all" — a family with no surface for one slot returns false + // rather than NULL'ing the hook. bool (*accepts_ext_kind)( const struct transcribe_model * model, transcribe_ext_slot slot, uint32_t kind); // Optional pre-clear validation for the one-shot transcribe_run path, - // the _RUN-slot analogue of stream_validate. The dispatcher calls it - // as the LAST gate before clear_result wipes the previous snapshot, - // and AFTER its own pre-clear checks: run_params struct size, the - // _RUN-slot ext header size + kind acceptance, and the run-param - // checks (enum range, timestamp ceiling, language, TRANSLATE support). - // Must be pure: zero mutation of ctx or model. On non-OK return the - // dispatcher returns the status with the previous result snapshot - // fully preserved, so a rejected run ext cannot destroy the prior - // transcript. - // - // Scope of the snapshot-preservation guarantee: a family SHOULD put - // its run-ext value validation here (the way parakeet's stream_validate - // vets its (L,C,R) menu) so a caller typo is caught pre-clear. The - // guarantee only extends to what this hook actually checks. A family - // that defers some value checks to its run() handler trades them out - // of the guarantee: a correctly-shaped but semantically-malformed ext - // can still be rejected post-clear. Whisper validates only the per-kind - // minimum ext size here (transcribe_ext_check against the full ext - // struct) and intentionally leaves its prompt-semantics checks in run() - // — an accepted gap on the one-shot path; see docs/follow-ups.md. + // the _RUN-slot analogue of stream_validate. Called as the LAST gate + // before clear_result wipes the snapshot, after the dispatcher's own + // pre-clear checks. Must be pure. On non-OK the prior snapshot is fully + // preserved, so a rejected run ext cannot destroy the prior transcript. // - // Placed last in the struct so families that do not accept a _RUN ext - // can leave it value-initialized to NULL; the dispatcher skips a NULL - // hook and the generic header-size + kind checks remain in force. + // A family SHOULD put its run-ext value validation here so a caller typo + // is caught pre-clear; the preservation guarantee only covers what this + // hook actually checks (a family that defers checks to run() trades them + // out of the guarantee). NULL is skipped; the generic header-size + kind + // checks remain in force. transcribe_status (*run_validate)( const struct transcribe_session * ctx, const transcribe_run_params * params); }; // Look up an architecture by name. Returns nullptr if no registered -// family matches. The lookup is O(n) over the registry; n is small -// (single-digit) and lookups happen once per load. +// family matches. const Arch * find_arch(const char * name); } // namespace transcribe diff --git a/src/transcribe-backend.h b/src/transcribe-backend.h index 1e4f9095..fda433b5 100644 --- a/src/transcribe-backend.h +++ b/src/transcribe-backend.h @@ -1,31 +1,15 @@ // transcribe-backend.h - internal backend selection types. // -// INTERNAL. C++17. Consumed only from other .cpp files inside src/. -// Not part of the public include/transcribe.h ABI. +// INTERNAL, C++17. Not part of the public ABI. // -// Rationale -// --------- -// The public API exposes a small `transcribe_backend_request` enum -// (auto|cpu|cpu_accel|metal|vulkan|cuda). Internally the library needs two related -// things that don't belong in the public header: -// -// 1. A typed classification of whatever ggml backend we actually -// landed on — `BackendKind`. This replaces string-matching on -// `ggml_backend_name()` scattered across model code (flash-attn -// policy, primary-backend-is-CPU policy, etc). -// -// 2. A small "backend plan" struct — `BackendPlan` — that holds the -// original request, the primary backend handle, its classified -// kind, and the full scheduler list in priority order. Every -// per-family load() resolves one of these, stores it on the -// model, and hands it to helpers that need to key off the -// primary backend (F16→F32 conv pointwise promotion, flash-attn -// default decisions). -// -// Keeping these out of transcribe-load-common.h means -// transcribe-flash-policy can include this header without dragging -// in the whole load scaffolding, and future consumers can ask -// "which backend is this?" without going through the loader. +// The public API exposes a small transcribe_backend_request enum. +// Internally the library needs two things that don't belong in the public +// header: a typed classification of the ggml backend it landed on +// (BackendKind, replacing string-matching on ggml_backend_name()), and a +// BackendPlan holding the request, primary backend handle, its kind, and +// the scheduler list in priority order. Every per-family load() resolves a +// plan and hands it to helpers that key off the primary backend (F16→F32 +// conv promotion, flash-attn defaults). #pragma once @@ -42,11 +26,9 @@ typedef struct ggml_backend_device * ggml_backend_dev_t; namespace transcribe { -// Classified backend kind. This is a library-internal classification -// based on the ggml backend device type plus backend registry name. -// -// Only the kinds the library actually needs to branch on are listed; -// anything else is Unknown. +// Library-internal classification from the ggml backend device type plus +// registry name. Only the kinds the library branches on are listed; anything +// else is Unknown. enum class BackendKind { Unknown = 0, Cpu, // ggml CPU backend (strict system memory) diff --git a/src/transcribe-batch-util.h b/src/transcribe-batch-util.h index f3a66d78..c9b84a96 100644 --- a/src/transcribe-batch-util.h +++ b/src/transcribe-batch-util.h @@ -1,14 +1,11 @@ // transcribe-batch-util.h - shared helpers for the offline batch path // (transcribe_run_batch family run_batch() hooks). // -// The per-utterance feature extraction (mel / kaldi-fbank / LFR) that every -// family runs at the top of run_batch() is pure host code with no -// cross-utterance state, so it is embarrassingly parallel across the batch. -// When the encoder runs on a fast accelerator the host feature extraction is -// frequently the dominant wall cost, so batching it across CPU workers is the -// single biggest end-to-end lever. Families differ in their frontend type and -// compute() signature, so the work itself is supplied as a per-index callable; -// this header owns only the parallel-dispatch boilerplate. +// Per-utterance feature extraction (mel / kaldi-fbank / LFR) is pure host +// code with no cross-utterance state, so it is embarrassingly parallel +// across the batch — and often the dominant wall cost when the encoder +// runs on a fast accelerator. Families supply the per-index work as a +// callable; this header owns the parallel-dispatch boilerplate. #pragma once @@ -28,56 +25,51 @@ namespace transcribe { // Invoke `work(i)` for every i in [0, n) across up to `n_threads` worker // threads (clamped to [1, n]; n_threads <= 0 means hardware_concurrency()). -// The caller's `work` MUST be reentrant: each index runs on exactly one -// thread, with no shared mutable state between indices (write per-index -// outputs into caller-owned, index-disjoint storage). Returns true iff every -// invocation returned true; a false return from any index does not stop the -// others (the caller decides what to do with a partial failure — typically -// fall back to the per-utterance path for the whole call). +// `work` MUST be reentrant: each index runs on one thread with no shared +// mutable state (write per-index outputs into index-disjoint storage). +// Returns true iff every invocation returned true; a false return does not +// stop the other indices (the caller handles partial failure — typically by +// falling back to the per-utterance path). bool parallel_for_all(int n, int n_threads, const std::function & work); -// Default CPU thread count for the ggml compute backends and the host -// parallel-for, for when the caller passes n_threads <= 0. Counts the CPUs the -// process may actually run on (the affinity mask via sched_getaffinity on Linux -// / GetProcessAffinityMask on Windows), NOT the host's total core count, then -// clamps to [1, cap]. This matters under a constrained scheduler (taskset, a -// cpuset, a CI runner pinned to N vCPUs): std::thread::hardware_concurrency() -// reports the full host, so the old `min(cap, hardware_concurrency())` default -// oversubscribed the usable cores and ggml's spin-wait barriers then livelocked. -// CAVEAT: the affinity mask reflects taskset/cpuset, but NOT a CFS bandwidth +// Default CPU thread count for the ggml backends and the host parallel-for +// when the caller passes n_threads <= 0. Counts the CPUs the process may +// actually run on (the affinity mask via sched_getaffinity on Linux / +// GetProcessAffinityMask on Windows), NOT the host's total core count, then +// clamps to [1, cap]. This matters under a constrained scheduler (taskset, +// cpuset, a CI runner pinned to N vCPUs): hardware_concurrency() reports the +// full host, so a `min(cap, hardware_concurrency())` default oversubscribes +// the usable cores and ggml's spin-wait barriers livelock. +// CAVEAT: the affinity mask reflects taskset/cpuset but NOT a CFS bandwidth // quota (`docker --cpus=N` without --cpuset) — that case still over-counts. // A cap <= 0 disables the clamp (use all usable CPUs). int default_n_threads(int cap = 8); // Resolve a CPU thread count and apply it to every CPU/BLAS backend in `sched` -// via the backend registry's ggml_backend_set_n_threads. `requested <= 0` means -// "use default_n_threads()". GPU backends don't expose the setter and are -// skipped; a null sched is a no-op. Returns the resolved thread count so callers -// can reuse it (e.g. for the host parallel-for). This is the ONE place archs -// configure scheduler threads — call it instead of hand-rolling the backend -// loop, so the affinity-aware default can never be forgotten at a new site. +// via ggml_backend_set_n_threads. `requested <= 0` means default_n_threads(). +// GPU backends don't expose the setter and are skipped; a null sched is a +// no-op. Returns the resolved count for reuse (e.g. the host parallel-for). +// The ONE place archs configure scheduler threads — call it instead of hand- +// rolling the loop so the affinity-aware default is never forgotten. int configure_sched_n_threads(ggml_backend_sched_t sched, int requested); // --------------------------------------------------------------------------- // Encoder-batch scaffolding (the run_batch() fast-path recipe) // -// Every batched-encoder family runs the same handful of host-side steps around -// its (family-specific) encoder graph: pack+pad each utterance's features into -// one batch slab, build per-utterance padding masks, then host-slice the shared -// encoder output and decode each utterance. These helpers own the parts that -// are byte-for-byte identical across families so a new family wires the recipe -// instead of re-deriving it. They produce exactly the host buffers the -// hand-written paths did (numerically identical), and no-op cleanly on a null -// tensor so a caller can pass a mask that the graph did not allocate. +// Every batched-encoder family runs the same host-side steps around its +// encoder graph: pack+pad each utterance's features into one batch slab, build +// per-utterance padding masks, then host-slice the shared encoder output and +// decode each utterance. These helpers own the parts that are identical across +// families. They no-op cleanly on a null tensor so a caller can pass a mask the +// graph did not allocate. // --------------------------------------------------------------------------- // Pack per-utterance channel-major features into one batched slab and zero-pad // each along time. `src[b]` is channel-major [n_ch, lens[b]] (element (c,t) at -// src[c*lens[b] + t]); `dst` is sized n_ch * T_max * n and laid out so slab b -// is [T_max, n_ch] (element (t,c) at (b*n_ch + c)*T_max + t), matching the -// conformer mel_in tensor ne = [T_max, n_ch, 1, B]. The padded tail -// (t >= lens[b]) is left zero. +// src[c*lens[b] + t]); `dst` is sized n_ch * T_max * n with slab b laid out as +// [T_max, n_ch] (element (t,c) at (b*n_ch + c)*T_max + t), matching the +// conformer mel_in tensor ne = [T_max, n_ch, 1, B]. Padded tail left zero. void pack_pad_channel_major( std::vector & dst, const std::vector> & src, @@ -105,13 +97,11 @@ void fill_valid_frame_mask(ggml_tensor * mask, // Per-utterance decode loop for the batched-encoder fast path. For each // utterance b in [0, n): polls abort, resets the scratch result slot, then -// calls `decode_fn(b, host_buf + b*utt_elems)` to run the family's host decode -// into that slot. Each result is snapshotted into session->batch_results with -// its mel/encode timing set to the per-utterance amortization of the single -// shared encode + total mel cost (so the per-utt timings sum to the real batch -// wall cost; decode time is whatever decode_fn recorded). `decode_fn`'s return -// is recorded as that utterance's status. Returns TRANSCRIBE_ERR_ABORTED if -// abort fires between utterances, else TRANSCRIBE_OK. +// calls `decode_fn(b, host_buf + b*utt_elems)`. Each result is snapshotted into +// session->batch_results with mel/encode timing set to the per-utterance +// amortization of the shared encode + total mel cost (so per-utt timings sum to +// the real batch wall cost); decode_fn's return is the utterance's status. +// Returns TRANSCRIBE_ERR_ABORTED if abort fires between utterances, else OK. transcribe_status decode_batch_slices( transcribe_session * session, int n, @@ -124,14 +114,13 @@ transcribe_status decode_batch_slices( // --------------------------------------------------------------------------- // Batched encoder-decoder greedy step loop (cohere / canary / moonshine) // -// The greedy enc-dec families share one decode driver: a uniform prompt feed of -// `prompt_len` lockstep steps over `prompt_ids`, then greedy generation until -// every row hits eos_id / max_new / the KV window. Self-attention keys are -// masked per row as positions fill; the KV read-window grows on demand -// (rebuilding the step graph) so short batches don't pay for the full n_ctx. -// moonshine is the degenerate no-growth case (init_window == max_n_kv, a single -// decoder_start prompt token). The cross-attention mask is static across steps -// and is uploaded by the family's rebuild callback (it owns the cross-KV cache). +// Shared decode driver: a uniform prompt feed of `prompt_len` lockstep steps +// over `prompt_ids`, then greedy generation until every row hits eos_id / +// max_new / the KV window. Self-attention keys are masked per row as positions +// fill; the KV read-window grows on demand (rebuilding the step graph) so short +// batches don't pay for the full n_ctx. moonshine is the degenerate no-growth +// case (init_window == max_n_kv, single decoder_start token). The cross- +// attention mask is static and uploaded by the family's rebuild callback. // --------------------------------------------------------------------------- // The B-utterance batched step graph as built by a family's @@ -162,11 +151,9 @@ using EncDecRebuildFn = std::function; // *n_steps_out (if non-null) receives the number of compute steps run. // // truncated_out (if non-null) is sized to n_batch and set per row: 1 when that -// (valid) row stopped because the loop hit the generation budget (max_new) or -// the context window (max_n_kv) BEFORE emitting eos_id — i.e. its transcript -// was truncated — and 0 otherwise (reached eos, or invalid). This lets a -// family report per-utterance TRANSCRIBE_ERR_OUTPUT_TRUNCATED from run_batch, -// matching the single-utterance contract. See docs/input-limits.md. +// (valid) row hit the generation budget (max_new) or the context window +// (max_n_kv) BEFORE emitting eos_id (transcript truncated), else 0. Lets a +// family report per-utterance TRANSCRIBE_ERR_OUTPUT_TRUNCATED from run_batch. transcribe_status run_batched_encdec_step_loop( transcribe_session * session, ggml_backend_sched_t sched, diff --git a/src/transcribe-debug.h b/src/transcribe-debug.h index 5603dada..f10befc2 100644 --- a/src/transcribe-debug.h +++ b/src/transcribe-debug.h @@ -1,46 +1,24 @@ // transcribe-debug.h - per-stage tensor dumping for the numerical // accuracy harness. // -// Phase 4 step 2 of the encoder port. Provides a single internal API -// the encoder graph builder can call to write a named tensor's -// contents to disk in the same on-disk format that -// scripts/dump_reference_*.py emits, so scripts/compare_tensors.py can -// diff a C++ dump dir against a Python reference dump dir without -// translation. -// -// On-disk format (matches scripts/dump_reference_*.py::write_dump): +// Writes a named tensor to disk in the format scripts/dump_reference_*.py +// emits, so scripts/compare_tensors.py can diff C++ vs Python dump dirs +// directly. // +// On-disk format (matches dump_reference_*.py::write_dump): // /.f32 raw little-endian fp32, row-major (C order) -// /.json sidecar metadata -// -// The JSON sidecar carries: -// -// { "name": "...", -// "stage": "...", // optional provenance label -// "shape": [d0, d1, ...], // numpy/row-major (slow-to-fast) -// "dtype": "f32", -// "layout": "row-major", -// "min": , "max": , "mean": , -// "source": { "kind": "cpp" } } -// -// Critical layout note: ggml ne[] is fast-to-slow (ne[0] = innermost, -// most-contiguous dim). Numpy / the Python dumpers use slow-to-fast -// (last axis varies fastest). The dumper converts by reversing ne[] -// and dropping trailing 1s, so a ggml tensor with ne=[1024, 275, 1, 1] -// (typical encoder activation [d_model, T, B=1]) lands on disk as -// shape=[275, 1024] — same as the reference's [T, d_model] tensor -// after squeezing the batch dim. -// -// Activation: dumping is gated on the TRANSCRIBE_DUMP_DIR environment -// variable. When unset, init() is a no-op and every dump_tensor call -// returns immediately. The encoder graph builder can sprinkle -// dump_tensor() calls freely without affecting normal-build cost. -// -// Thread safety: the dump dir is captured once at init() and the -// per-call file writes are not synchronized. The harness assumes -// single-threaded use during numerical accuracy bringup. Concurrent -// dumps would race on file creation and produce truncated output — -// not a goal for v1. +// /.json sidecar: { name, stage?, shape (slow-to-fast), +// dtype, layout, min, max, mean, source } +// +// Critical layout note: ggml ne[] is fast-to-slow (ne[0] = innermost). +// Numpy / the Python dumpers use slow-to-fast. The dumper reverses ne[] +// and drops trailing 1s, so ggml ne=[1024, 275, 1, 1] ([d_model, T, B=1]) +// lands on disk as shape=[275, 1024] — the reference's [T, d_model] after +// squeezing the batch dim. +// +// Gated on TRANSCRIBE_DUMP_DIR: unset → init() is a no-op and dumps return +// immediately. Not thread-safe (file writes are unsynchronized); single- +// threaded use only. #pragma once @@ -108,10 +86,8 @@ void dump_tensor(const char * name, // Dump a host-side fp32 buffer to /.{f32,json}. // -// Used by code paths that don't run on a ggml backend — phase 5's -// decoder runs the LSTM + joint + TDT loop directly on host floats -// because the per-step compute is small and a backend graph would -// add lifetime complexity for no perf win. +// Used by code paths that don't run on a ggml backend — e.g. a decoder +// that runs the LSTM + joint + TDT loop directly on host floats. // // `data` points at `n_elem` contiguous fp32 values. `shape` is the // numpy/row-major (slow-to-fast) shape and must satisfy diff --git a/src/transcribe-flash-policy.h b/src/transcribe-flash-policy.h index 02912851..0ebf2df5 100644 --- a/src/transcribe-flash-policy.h +++ b/src/transcribe-flash-policy.h @@ -1,46 +1,22 @@ // transcribe-flash-policy.h - shared flash-attention policy helpers. // -// INTERNAL. C++17. Consumed only from other .cpp files inside src/. -// Not part of the public include/transcribe.h ABI. +// INTERNAL, C++17. Not part of the public ABI. // -// Flash-attention support in ggml is per-backend and per-head-dim. -// Each family picks its own defaults based on encoder/decoder head -// dimensions and which backend it was loaded on. -// -// The "which backend" half of that decision now lives in -// transcribe-backend.h as BackendKind — families check -// `plan.primary_kind == BackendKind::Metal` directly instead of -// calling a helper here. -// -// What stays in this header is the one cross-family concern that is -// not backend-classification: honoring the TRANSCRIBE_NO_FLASH / -// TRANSCRIBE_FORCE_FLASH environment overrides. These are a global -// "no flash anywhere" / "flash everywhere" debug lever, not a -// per-family policy knob. -// -// The default on/off decision itself is still per-family: it -// depends on the family's head dim (cohere encoder=160 unsupported -// on Metal, cohere decoder=128 supported; parakeet encoder=128 -// supported). Extracting that decision here would paper over a real -// architectural difference. +// Flash-attention support in ggml is per-backend and per-head-dim, so the +// default on/off decision is per-family (it depends on head dim and the +// backend it was loaded on). What lives here is the one cross-family +// concern: honoring the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH global +// debug overrides. #pragma once namespace transcribe::flash { -// Apply the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH -// environment overrides to a pair of encoder/decoder flash flags. -// Either override forces both flags in the same direction -- the -// user's intent is "no flash kernels anywhere" or "flash kernels -// everywhere", not "flash on encoder but not decoder". -// -// If both env vars are set simultaneously, FORCE wins (last apply). -// That's a deliberately undocumented edge case; the interesting -// combination is neither-or-one, not both. -// -// encoder_use_flash: in/out. Updated in-place if an override is -// set; left untouched otherwise. -// decoder_use_flash: same. +// Apply the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH environment +// overrides to a pair of encoder/decoder flash flags. Either override forces +// both flags in the same direction; if both are set, FORCE wins. The flags +// are in/out: updated in-place if an override is set, left untouched +// otherwise. void apply_env_overrides(bool & encoder_use_flash, bool & decoder_use_flash); diff --git a/src/transcribe-load-common.cpp b/src/transcribe-load-common.cpp index fe4cf5ad..6d5f78a9 100644 --- a/src/transcribe-load-common.cpp +++ b/src/transcribe-load-common.cpp @@ -1,10 +1,7 @@ // transcribe-load-common.cpp - shared model-load scaffolding. // -// See transcribe-load-common.h for rationale. The functions here -// are the common backend-init and tensor-stream logic that used to -// live duplicated in every per-family model.cpp. The pre-hoist copies -// differed only in the "parakeet:"/"cohere:" log prefix and were -// otherwise character-identical. +// See transcribe-load-common.h for rationale. The functions here are the +// common backend-init and tensor-stream logic shared by per-family load(). #include "transcribe-load-common.h" @@ -250,18 +247,13 @@ transcribe_status init_backends(transcribe_backend_request requested, switch (requested_raw) { case TRANSCRIBE_BACKEND_CPU: case TRANSCRIBE_BACKEND_CPU_ACCEL: { - // CPU primary. The CPU_ACCEL variant additionally layers the - // host-memory accelerators (BLAS/AMX/…) registered by ggml as - // GGML_BACKEND_DEVICE_TYPE_ACCEL onto the scheduler. Both - // variants set primary_kind == Cpu so CPU-keyed policy - // decisions (e.g. F16→F32 conv pointwise promotion) trigger - // identically. Strict CPU is the right choice for numerical - // reference runs and cross-platform determinism; CPU_ACCEL is - // for production CPU throughput on machines where the build - // included an accel backend (e.g. GGML_BLAS=ON). + // CPU primary. CPU_ACCEL additionally layers the host-memory + // accelerators (BLAS/AMX/…, GGML_BACKEND_DEVICE_TYPE_ACCEL) onto the + // scheduler. Both set primary_kind == Cpu so CPU-keyed policy (e.g. + // F16→F32 conv pointwise promotion) triggers identically. // - // ggml requires the CPU backend to sit last in the scheduler - // list, so accel backends (when included) go in first. + // ggml requires the CPU backend to sit last in the scheduler list, + // so accel backends (when included) go in first. const bool with_accel = (requested == TRANSCRIBE_BACKEND_CPU_ACCEL); ggml_backend_t cpu_be = init_cpu_backend(error_tag); if (cpu_be == nullptr) return TRANSCRIBE_ERR_BACKEND; @@ -434,10 +426,9 @@ transcribe_status promote_conv_pw_f16_to_f32_on_cpu( ggml_context ** out_ctx, ggml_backend_buffer_t * out_buffer) { - // Key off the classified primary kind, not off ACCEL/CPU - // ordering in the backend list. This is the fix for the latent - // bug where strict-CPU requests could be silently shadowed by - // an ACCEL backend sitting ahead of CPU. + // Key off the classified primary kind, not off ACCEL/CPU ordering in + // the backend list, so a strict-CPU request reliably triggers promotion + // even when an ACCEL backend sorts ahead of CPU. if (plan.primary_kind != BackendKind::Cpu) return TRANSCRIBE_OK; if (plan.primary == nullptr) return TRANSCRIBE_OK; diff --git a/src/transcribe-load-common.h b/src/transcribe-load-common.h index c8ebf81b..c6fdd4ce 100644 --- a/src/transcribe-load-common.h +++ b/src/transcribe-load-common.h @@ -1,23 +1,18 @@ // transcribe-load-common.h - shared model-load scaffolding. // -// INTERNAL. C++17. Consumed only from other .cpp files inside src/. -// Not part of the public include/transcribe.h ABI. +// INTERNAL, C++17. Not part of the public ABI. // -// Every per-family load() walks the same three steps after the GGUF -// catalog has been built: +// Every per-family load() walks the same three steps after building the +// GGUF catalog: // // 1. Resolve a BackendPlan from transcribe_model_load_params::backend. -// (See transcribe-backend.h for the plan type.) -// 2. Allocate a backend buffer for every tensor in ctx_meta on the -// primary backend from the plan. -// 3. Stream each tensor's bytes from the GGUF data section into -// its backend slot via ggml_backend_tensor_set. +// 2. Allocate a backend buffer for every tensor in ctx_meta on the plan's +// primary backend. +// 3. Stream each tensor's bytes from the GGUF data section into its +// backend slot via ggml_backend_tensor_set. // -// Parakeet and Cohere had character-for-character copies of all -// three steps, modulo a "parakeet:"/"cohere:" log prefix. This -// header hoists them. Per-family `load()` still owns the model, -// params validation, weights catalog, fusion passes, and destructor -// wiring; what moves here is the uniform scaffolding in between. +// These functions hoist that scaffolding; per-family load() still owns the +// model, params validation, weights catalog, fusion, and destructor wiring. #pragma once @@ -158,13 +153,8 @@ transcribe_status promote_conv_pw_f16_to_f32_on_cpu( // // On success, `out` is resized and filled. On every other return // (including Absent), `out` is cleared so stale data cannot leak. -// On any failure other than Absent, a diagnostic is printed to -// stderr using `error_tag` and `tensor_name`. -// -// Introduced for the cohere frontend tensor path -// (frontend.mel_filterbank / frontend.window), which previously used -// an unchecked inline lambda that silently accepted non-F32 types, -// didn't verify read counts, and had no expected-size bounds. +// On any failure other than Absent, a diagnostic is logged using +// `error_tag` and `tensor_name`. enum class ReadF32Result { Ok, // tensor found, validated, and read successfully Absent, // tensor not in the GGUF index (caller should compute) diff --git a/src/transcribe-loader.h b/src/transcribe-loader.h index d768b26a..bc4ebf70 100644 --- a/src/transcribe-loader.h +++ b/src/transcribe-loader.h @@ -1,20 +1,11 @@ // transcribe-loader.h - internal GGUF loader (architecture-agnostic). // -// This header is INTERNAL. It is C++17 (not C) on purpose: it is consumed -// only from other .cpp files inside src/. The public C ABI lives in -// include/transcribe.h and never references anything declared here. -// -// The Loader exists so transcribe_model_load_file can answer two -// questions before it has to commit to a per-family code path: -// -// 1. Is this file a readable GGUF at all? -// 2. What architecture does it claim (general.architecture / stt.variant)? -// -// Once the architecture is known the loader is handed off to the per-arch -// handler, which may take ownership of the underlying gguf_context via -// release_gguf(). If the handler does not take ownership (e.g. because it -// returned NOT_IMPLEMENTED in pass 2A), the Loader's destructor frees the -// context normally on stack unwinding. +// INTERNAL, C++17. The Loader lets transcribe_model_load_file answer two +// questions before committing to a per-family code path: is this a readable +// GGUF, and what architecture does it claim (general.architecture / +// stt.variant)? It is then handed to the per-arch handler, which may take +// ownership of the gguf_context via release_gguf(); if it does not, the +// Loader's destructor frees the context on stack unwinding. #pragma once @@ -32,9 +23,8 @@ class Loader { Loader() = default; ~Loader(); - // Non-copyable. Move semantics are not needed in 0.x because the - // loader is always either constructed-then-released-by-handler or - // constructed-then-destroyed-on-the-stack. + // Non-copyable; move is not needed (the loader is always either + // released to the handler or destroyed on the stack). Loader(const Loader &) = delete; Loader & operator=(const Loader &) = delete; Loader(Loader &&) = delete; diff --git a/src/transcribe-log.h b/src/transcribe-log.h index 137b55c8..ebca527e 100644 --- a/src/transcribe-log.h +++ b/src/transcribe-log.h @@ -1,20 +1,12 @@ // transcribe-log.h - internal printf-style logging facility. // // Routes a formatted message through the process-global log callback -// installed via transcribe_log_set(), falling back to stderr when no -// callback is installed (so dev / CLI builds still see the diagnostic). -// -// This is the supported way for library internals — including per-family -// run() drivers in src/arch// — to surface diagnostics. Family -// code must NOT call std::fprintf(stderr, ...) directly for user-facing -// conditions: a consumer that installed a log sink would never see those -// messages. In particular the input-length contract (input-too-long -// rejection, mid-decode truncation, soft-window degradation; see -// docs/input-limits.md) is surfaced through here so it reaches the -// caller's sink at the right level. -// -// The implementation lives in transcribe.cpp alongside the callback -// globals it reads. +// installed via transcribe_log_set(), falling back to stderr when none is +// installed. This is the supported way for library internals (including +// per-family run() drivers) to surface diagnostics; family code must NOT +// call std::fprintf(stderr, ...) directly, or a consumer's log sink would +// never see the message. The implementation lives in transcribe.cpp +// alongside the callback globals it reads. #pragma once diff --git a/src/transcribe-mel.cpp b/src/transcribe-mel.cpp index 1933d38a..111438fb 100644 --- a/src/transcribe-mel.cpp +++ b/src/transcribe-mel.cpp @@ -1,63 +1,27 @@ // transcribe-mel.cpp - native C++ log-mel feature extractor. // -// Implementation notes -// -------------------- +// Matches NeMo's AudioToMelSpectrogramPreprocessor +// (FilterbankFeatures.forward in features.py) as exported in the +// Parakeet 0.6B nemo128.onnx preprocessor. Constants / ordering: // -// The pipeline below matches NeMo's AudioToMelSpectrogramPreprocessor -// (FilterbankFeatures.forward in nemo/collections/asr/parts/preprocessing/ -// features.py) as exported in the Parakeet 0.6B nemo128.onnx -// preprocessor. Phase 3 preflight cross-checked the algorithm three -// ways: NeMo source, ONNX op graph + initializers, and a standalone -// PyTorch reproduction. Constants and ordering decisions: -// -// * Pre-emphasis alpha: read from MelConfig::pre_emphasis (0.97 for -// real Parakeet; 0.0 disables the step entirely). -// -// * Reflect padding by n_fft/2 on each side. The NeMo source on -// main says pad_mode="constant" but the ONNX export uses -// reflect, and the released checkpoint was packaged with the -// reflect path; reflect is what produces working transcriptions. -// Trust the ONNX, ignore the source on this one point. -// -// * STFT in fp32 for the mixed-radix path (whisper, n_fft=400) and +// * Pre-emphasis alpha: MelConfig::pre_emphasis (0.97; 0.0 disables). +// * Reflect padding by n_fft/2 on each side. The NeMo source says +// pad_mode="constant" but the ONNX export (and the released +// checkpoint) use reflect; trust the ONNX. +// * STFT in fp32 for the mixed-radix path (whisper, n_fft=400), // fp64 for the radix-2 vDSP path (parakeet/cohere, n_fft=512). -// The fp32 mixed-radix matches whisper.cpp's choice and was -// verified WER-neutral on test-clean-300 (5.83% vs 5.84% fp64). -// vDSP's fp64 stays because it's already fast (Accelerate SIMD) -// and the precision is free. -// -// * Window: symmetric Hann length win_length, zero-padded to n_fft. -// The NeMo source explicitly passes periodic=False -// (features.py:316). -// -// * Power spectrum: re^2 + im^2 (matches ONNX ReduceSumSquare). -// mag_power=2.0 is collapsed into the squaring; we never compute -// the magnitude itself. -// -// * Mel filterbank: Slaney-normalized librosa formula computed at -// constructor time from (sr, n_fft, n_mels, fmin, fmax). Verified -// to match the ONNX-baked filterbank to ~3e-7 in fp32. -// -// * Log epsilon: 2^-24 (NeMo log_zero_guard_value default). -// Hardcoded; not in the schema. Different from the 1e-5 we -// initially planned for; the value comes straight from -// features.py:256. -// -// * Per-feature normalize uses the UNBIASED variance estimator, -// dividing by (n_frames - 1) not n_frames. This is what NeMo's -// normalize_batch does (features.py line 79) and what the ONNX -// graph reproduces (Sub-1 op before the Div). The biased -// estimator would silently drift the test by ~1/n_frames. +// fp32 mixed-radix matches whisper.cpp and is WER-neutral. +// * Window: symmetric Hann length win_length zero-padded to n_fft +// (NeMo passes periodic=False, features.py:316). +// * Power spectrum: re^2 + im^2 (mag_power=2.0 folded into squaring). +// * Mel filterbank: Slaney-normalized librosa formula, ~3e-7 vs ONNX. +// * Log epsilon: 2^-24 (NeMo log_zero_guard_value, features.py:256). +// * Per-feature normalize: UNBIASED variance, divides by (n-1) +// (NeMo normalize_batch, features.py:79); biased drifts ~1/n. +// * Normalize epsilon: 1e-5 (NeMo CONSTANT). +// * Dither: never applied (NeMo gates it on self.training). // -// * Normalize epsilon: 1e-5. Matches NeMo's CONSTANT. -// -// * Dither: never applied. NeMo gates dither on self.training, -// so inference always sees a deterministic frontend. We don't -// even read MelConfig::dither. -// -// Output layout: [num_mels, n_frames] row-major float32. The encoder -// builder in phase 4 will copy this into a ggml_tensor at the -// backend boundary; the frontend itself stays backend-free. +// Output: [num_mels, n_frames] row-major float32; backend-free. #include "transcribe-mel.h" @@ -197,18 +161,9 @@ void build_hann_window_symmetric_padded( // surface zero. Drop in pocketfft only if profiling later shows the // frontend is a bottleneck. // Per-MelFrontend sin/cos cache, indexed by 2π*i / lut_size. The mixed- -// radix path is the only consumer; for it to hit the LUT on every level -// we need lut_size divisible by every N seen during recursion. With -// lut_size == n_fft this holds automatically: n_fft → n_fft/2 → … → odd -// leaf, all divisors of n_fft. Stride at recursion level N is -// lut_size / N, and (k*n*stride) % lut_size lands on the right phase. -// -// Historical note: this was previously a process-wide 5040-entry -// singleton ("7! covers divisors 1..10"), but for Whisper's n_fft=400 = -// 2^4·5² the recursion sizes are {400, 200, 100, 50, 25} — none divide -// 5040, so use_lut was false at every level and we silently fell -// through to live std::cos/sin in tight loops (~10K trig calls per -// frame from the leaf DFT alone). Sizing the LUT to n_fft fixes that. +// radix path is the only consumer; lut_size == n_fft so every recursion +// size N (n_fft → n_fft/2 → … → odd leaf) divides it. Stride at level N +// is lut_size / N, and (k*n*stride) % lut_size lands on the right phase. // Naive O(N^2) DFT leaf for the mixed-radix path. Used when N hits an // odd factor during recursive halving (e.g. N=400 → 200 → 100 → 50 → 25 @@ -431,15 +386,10 @@ transcribe_status MelFrontend::compute( // ---- 1. Pad + pre-emphasis ---- // - // The non-pow2 (mixed-radix) path runs fp32 throughout, so we build - // the padded buffer directly in fp32 and skip the intermediate emph - // copy entirely. The pow2 paths (vDSP / hand-rolled radix-2) stay - // fp64 since their FFT routines are fp64-native. - // - // Whisper's mel was already lossy-cast fp64→fp32 right before the - // FFT in the non-pow2 path, so this just removes a memory-bandwidth - // hit (and one allocation) without changing the precision the FFT - // itself sees. + // The non-pow2 (mixed-radix) path runs fp32 throughout, so build the + // padded buffer directly in fp32 and skip the intermediate emph copy. + // The pow2 paths (vDSP / hand-rolled radix-2) stay fp64 since their + // FFT routines are fp64-native. const bool n_fft_is_pow2 = ((n_fft > 0) && ((n_fft & (n_fft - 1)) == 0)); std::vector padded_f32; @@ -536,30 +486,21 @@ transcribe_status MelFrontend::compute( // ---- 3. STFT + mel filterbank matmul + log ---- // - // Output: log_mel[n_mels, n_frames] row-major float32. + // Output: log_mel[n_mels, n_frames] row-major float32. Two paths by + // n_fft factorization: // - // Two distinct architectures depending on n_fft factorization: + // non-pow2 (Whisper, Qwen3-ASR): FFT worker fuses the per-frame mel + // matmul + log directly into log_mel (no shared power[]), matmul + // parallelized across stft_threads. Same as whisper.cpp; a win on + // no-BLAS CPUs. + // pow2 (Parakeet, Cohere): FFT worker writes shared power[], then a + // post-pass does matmul + log via cblas_sgemm (much faster than a + // per-frame scalar matmul). vDSP / fft_radix2 are fp64-native, so + // the fp64 power[] avoids cast traffic. // - // non-pow2 (Whisper, Qwen3-ASR): the FFT worker also runs the - // per-frame mel matmul + log directly into log_mel. No shared - // power[] buffer, and the matmul parallelizes across all - // stft_threads (it's per-frame work). This is the same - // architecture whisper.cpp uses, and on no-BLAS CPUs it's a - // significant win — the matmul cost moves from a single- - // threaded post-pass to fully threaded inline work. - // - // pow2 (Parakeet, Cohere): the FFT worker writes a shared - // power[] buffer, then a post-pass does the matmul + log. On - // Apple/Linux+OpenBLAS the post-pass uses cblas_sgemm, which - // is dramatically faster than a per-frame scalar matmul, so - // worker fusion would be a regression there. The fp64 vDSP / - // fft_radix2 routines are also fp64-native, so keeping the - // existing fp64 power[] avoids cast traffic at the worker. - // - // log mode (whisper_mode): per_utterance writes log10(max(x, 1e-10)) - // so the normalize step skips the 1/ln(10) conversion. per_feature - // writes log(x + 2^-24) (natural log; base doesn't matter for the - // per-bin mean/std normalize that follows). + // log mode (whisper_mode): per_utterance writes log10(max(x, 1e-10)); + // per_feature writes log(x + 2^-24) (natural log; base is irrelevant + // to the per-bin mean/std normalize that follows). std::vector log_mel( static_cast(n_mels) * static_cast(n_frames)); @@ -593,14 +534,9 @@ transcribe_status MelFrontend::compute( if (!n_fft_is_pow2) { // Fused worker: window mul → mixed-radix FFT → power → mel - // matmul + log → log_mel, all in fp32, all per-frame on this - // thread. Per-thread scratches stay tiny (~14 KB total) and - // L1-resident; nothing crosses thread boundaries except the - // final log_mel writes. - // - // Matmul: outer m (mel), inner k (freq) with 4-way unroll. - // fp64 sum accumulator for numerical stability vs the linear - // sum size; cast back to fp32 at the log boundary. + // matmul + log → log_mel, all fp32, all per-frame on this thread. + // Matmul uses a 4-way-unrolled fp64 accumulator for numerical + // stability; cast back to fp32 at the log boundary. auto worker = [&](int tid) { std::vector fft_in(2 * static_cast(n_fft), 0.0f); std::vector fft_out(8 * static_cast(n_fft), 0.0f); @@ -782,39 +718,33 @@ transcribe_status MelFrontend::compute( } } #endif - // power[] released at end of this scope. - } // end pow2 path + } std::vector().swap(padded); std::vector().swap(padded_f32); std::vector().swap(window_f32); // ---- 5n. No-op normalize (NeMo "NA"/none) ---- - // Emit raw log-mel as-is. The feature normalisation that streaming - // Conformer variants like nemotron-speech-streaming-en-0.6b apply - // is baked into training (mean/std absorbed into the encoder - // weights), so at inference we hand the encoder unnormalised - // log-mel exactly the way NeMo's `normalize_batch` falls through. + // Emit raw log-mel as-is: streaming Conformer variants (e.g. + // nemotron-speech-streaming-en-0.6b) bake feature normalization into + // the encoder weights, matching NeMo's normalize_batch fall-through. // - // NeMo's FilterbankFeatures.forward still masks frames beyond - // `get_seq_len(n_samples) = floor(n_samples / hop_length)` to - // `pad_value` (zero) — that runs after normalize regardless of - // normalize_type. The reflect-pad STFT here emits one extra frame - // (n_frames = floor(n_samples / hop) + 1), so the last frame is a - // padding artifact and we mirror NeMo by zeroing it. Without this - // the encoder sees a real log-mel value at the padding position - // and the resulting frame-level drift dominates the early + // Trailing-frame masking (the canonical rationale): NeMo's + // FilterbankFeatures.forward masks frames beyond get_seq_len = + // floor(n_samples / hop_length) to pad_value (zero), after normalize, + // regardless of normalize_type. The reflect-pad STFT emits one extra + // frame (n_frames = floor(n/hop) + 1), so the last frame is a padding + // artifact and we zero it to mirror NeMo. Without this the encoder + // sees a real log-mel value there and the drift dominates the early // pre_encode output (max_abs ~13 on that single frame). if (cfg_.normalize == "none") { out_mel = std::move(log_mel); out_n_mels = n_mels; out_n_frames = n_frames; // pad_mode="none" emits exactly (n - win)/hop + 1 frames with no - // padding artifact, so the trailing-mask step is unnecessary - // (and would zero a legitimate final frame). For NeMo-style - // reflect/constant padding, the last centered frame is a padding - // artifact that needs to be zeroed to match NeMo's normalize_batch - // behavior — see the comment block in the doc-history above. + // padding artifact, so skip the trailing mask (it would zero a + // legitimate final frame). NeMo-style reflect/constant padding + // needs the mask — see the trailing-frame rationale above. if (!no_pad) { const int valid = static_cast( n_samples / static_cast(cfg_.hop_length)); @@ -830,10 +760,9 @@ transcribe_status MelFrontend::compute( } // ---- 5a. Whisper-style per-utterance normalization ---- - // log_mel already holds log10(max(x, 1e-10)) values (the matmul step - // wrote them directly). Find the global max, clamp to max - 8.0, - // then scale (x + 4) / 4. Whisper drops the trailing center-pad - // STFT frame so the output has n_samples / hop_length frames. + // log_mel already holds log10(max(x, 1e-10)). Find the global max, + // clamp to max - 8.0, then scale (x + 4) / 4. Drops the trailing + // center-pad STFT frame (output has n_samples / hop_length frames). if (cfg_.normalize == "per_utterance") { const int n_out = n_frames - 1; if (n_out <= 0) return TRANSCRIBE_ERR_INVALID_ARG; @@ -868,9 +797,9 @@ transcribe_status MelFrontend::compute( } // ---- 5b. Voxtral Realtime streaming log-mel (fixed global max) ---- - // Identical to per_utterance but the clamp floor uses a FIXED maximum - // (global_log_mel_max), not the per-utterance max — what makes per-frame - // normalization causal/streaming-safe. Drops the trailing center-pad frame. + // Like per_utterance but the clamp floor uses a FIXED max + // (global_log_mel_max), making per-frame normalization causal/ + // streaming-safe. Drops the trailing center-pad frame. if (cfg_.normalize == "global") { const int n_out = n_frames - 1; if (n_out <= 0) return TRANSCRIBE_ERR_INVALID_ARG; @@ -893,32 +822,24 @@ transcribe_status MelFrontend::compute( } // ---- 5. Per-feature normalize, unbiased ---- - // For each mel bin m: subtract the mean across time, divide by - // (std + 1e-5). The variance uses (n_norm - 1) in the - // denominator (Bessel's correction) to match NeMo's - // normalize_batch. fp64 accumulators throughout to keep the row - // sums numerically stable; cast back to fp32 at storage. + // For each mel bin m: subtract the time mean, divide by (std + 1e-5). + // Variance uses (n_norm - 1) (Bessel's correction) to match NeMo's + // normalize_batch. fp64 accumulators throughout, cast to fp32 at + // storage. // - // When pad_mode is "constant" (Cohere), the last STFT frame is - // a padding artifact. The normalization statistics are computed - // over the first seq_len = n_samples / hop frames (not n_frames - // = seq_len + 1), and the last frame is zeroed after - // normalization. This matches the Cohere reference - // (processing_cohere_asr.py get_seq_len / normalize_batch). + // pad_mode "constant" (Cohere): the last STFT frame is a padding + // artifact; stats are over the first seq_len = n/hop frames and the + // last is zeroed after normalize (processing_cohere_asr.py + // get_seq_len / normalize_batch). // - // When pad_mode is "reflect" (Parakeet / Canary), all n_frames - // frames are valid and used for normalization. Empirically, - // unconditionally treating the trailing frame as a center-pad - // artifact regresses long-form Canary decode (e.g. test-clean - // 672-122797-0008, 30.8s, drops into a "the boy was a man of - // the world" loop). NeMo's FilterbankFeatures masks per seq_len - // computed from the raw sample count, but for reflect-padded - // single-segment offline runs that seq_len == n_frames so no - // masking happens — keep the pad_mode gate to preserve that. + // pad_mode "reflect" (Parakeet / Canary): all n_frames are valid. + // Unconditionally masking the trailing frame regresses long-form + // Canary decode (test-clean 672-122797-0008, 30.8s, falls into a "the + // boy was a man of the world" loop): for reflect-padded single-segment + // offline runs NeMo's seq_len == n_frames so it masks nothing — the + // pad_mode gate preserves that. // - // resize() instead of assign(): the loop below writes every - // element exactly once, so the zero-fill that assign() would do - // is dead work. + // resize() not assign(): the loop writes every element once. const bool mask_last = (cfg_.pad_mode == "constant"); const int n_norm = mask_last ? (n_frames - 1) : n_frames; diff --git a/src/transcribe-mel.h b/src/transcribe-mel.h index f27d6711..d310eabb 100644 --- a/src/transcribe-mel.h +++ b/src/transcribe-mel.h @@ -1,30 +1,19 @@ // transcribe-mel.h - native C++ log-mel feature extractor. // -// PRIVATE header. Lives under src/ and is not exported. The public C -// ABI in include/transcribe.h does not surface this; the family -// handlers (currently only Parakeet) construct one and call it from -// inside transcribe_run when phase 4 lands. +// PRIVATE header (src/, not exported). Pure CPU, no ggml dependency: +// audio in, std::vector mel out; the family encoder builder +// copies the result into a ggml_tensor at the backend boundary. // -// Phase 3 scope: pure CPU, no ggml dependency. Audio in, mel out, -// std::vector for the buffer. The encoder builder in phase 4 -// is responsible for copying the result into a ggml_tensor at the -// backend boundary. +// Matches NeMo's AudioToMelSpectrogramPreprocessor / FilterbankFeatures +// as exported in nemo128.onnx. The per-NeMo numerical constants +// (log_eps = 2^-24, unbiased variance, reflect padding, fp64 STFT, +// mag_power = 2.0) are hardcoded in the .cpp rather than in the GGUF +// schema; they become per-family knobs only when a second family needs +// different values. // -// Algorithm pinned by phase 3 preflight (see RESUME.md "Parakeet -// frontend spec"). The implementation matches NeMo's -// AudioToMelSpectrogramPreprocessor / FilterbankFeatures pipeline as -// exported in nemo128.onnx, verified against the upstream NeMo -// features.py source and a standalone PyTorch reproduction. The -// per-NeMo numerical constants (log_eps = 2^-24, unbiased variance, -// reflect padding, fp64 STFT, mag_power = 2.0) are hardcoded in the -// .cpp file rather than added to the GGUF schema; see PLAN.md -// "complete contract" rule. They become per-family knobs only when a -// second family wants different values. -// -// Construction is one-shot: the constructor precomputes the periodic -// hann window (zero-padded from win_length to n_fft) and the -// Slaney-normalized mel filterbank. compute() is then a function of -// (config, audio) only. +// Construction is one-shot: the constructor precomputes the hann window +// (zero-padded from win_length to n_fft) and the Slaney-normalized mel +// filterbank. compute() is then a function of (config, audio) only. #pragma once diff --git a/src/transcribe-meta.cpp b/src/transcribe-meta.cpp index 9dece306..5eaf110c 100644 --- a/src/transcribe-meta.cpp +++ b/src/transcribe-meta.cpp @@ -204,7 +204,7 @@ transcribe_status read_required_u32_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } transcribe_status read_required_f32_kv(const gguf_context * gguf, @@ -224,7 +224,7 @@ transcribe_status read_required_f32_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } transcribe_status read_required_string_kv(const gguf_context * gguf, @@ -244,7 +244,7 @@ transcribe_status read_required_string_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } transcribe_status read_optional_bool_kv(const gguf_context * gguf, @@ -267,7 +267,7 @@ transcribe_status read_optional_bool_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } transcribe_status read_optional_string_kv(const gguf_context * gguf, @@ -290,7 +290,7 @@ transcribe_status read_optional_string_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } transcribe_status read_optional_int32_kv(const gguf_context * gguf, @@ -313,7 +313,7 @@ transcribe_status read_optional_int32_kv(const gguf_context * gguf, error_tag, key); return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable + return TRANSCRIBE_ERR_GGUF; } // --------------------------------------------------------------------------- @@ -335,7 +335,7 @@ transcribe_status read_capability_bool(const gguf_context * gguf, case KvResult::Ok: return TRANSCRIBE_OK; case KvResult::BadType: return TRANSCRIBE_ERR_GGUF; } - return TRANSCRIBE_ERR_GGUF; // unreachable; placates -Wreturn-type on some compilers + return TRANSCRIBE_ERR_GGUF; // unreachable; placates -Wreturn-type } } // namespace @@ -347,20 +347,10 @@ transcribe_status read_capability_kv(const gguf_context * gguf, return TRANSCRIBE_ERR_INVALID_ARG; } - // 2B-recognized capability keys. The schema is in PLAN.md - // "Speech-specific metadata" / `stt.capability.*`. - // - // Timestamp granularity (`max_timestamp_kind`) is deliberately - // NOT read here. Both shipped families have a fixed, code-set - // ceiling (Parakeet=TOKEN, Cohere=NONE) and no converter emits - // `stt.capability.timestamps`. A KV-driven override is a - // future change — the loader would need rules for how the KV - // interacts with what the family code can actually produce - // (the KV should only lower the ceiling, never raise it, since - // code is the floor of what can physically be emitted). When a - // second checkpoint of an existing family ships with a - // different ceiling, resolve that interaction before wiring - // the reader. + // Timestamp granularity (max_timestamp_kind) is deliberately NOT read + // here: every family has a fixed, code-set ceiling and no converter + // emits stt.capability.timestamps. A KV-driven override (which could + // only lower the ceiling, never raise it) is a future change. if (auto st = read_capability_bool(gguf, "stt.capability.translate", caps.supports_translate); st != TRANSCRIBE_OK) @@ -406,11 +396,8 @@ transcribe_status read_languages_kv(const gguf_context * gguf, case KvResult::BadType: return TRANSCRIBE_ERR_GGUF; } - - // Translation-target set: the target-side twin of general.languages, - // installed the same way (model-owned storage, republished caps - // pointer). Absent on ASR-only models and on GGUFs predating the KV, - // which leaves caps.n_translate_target_languages = 0 ("not advertised"). + // Optional translation metadata. The model stores copies so public caps + // pointers remain valid; absent keys leave old GGUFs permissive. std::vector targets; switch (read_string_array_kv(gguf, "stt.translation.target_languages", targets)) { case KvResult::Absent: @@ -422,9 +409,6 @@ transcribe_status read_languages_kv(const gguf_context * gguf, return TRANSCRIBE_ERR_GGUF; } - // Optional translation-pair contract. When present, the generic validator - // can reject explicitly supplied unsupported source>target pairs before - // family dispatch. Absent keeps old GGUFs permissive. std::vector pairs; switch (read_string_array_kv(gguf, "stt.translation.pairs", pairs)) { case KvResult::Absent: diff --git a/src/transcribe-meta.h b/src/transcribe-meta.h index dc282723..d8d1e589 100644 --- a/src/transcribe-meta.h +++ b/src/transcribe-meta.h @@ -1,38 +1,19 @@ // transcribe-meta.h - shared GGUF KV reading helpers + post-load // metadata reads (capabilities, languages). // -// This header is INTERNAL. The public C ABI in include/transcribe.h -// knows nothing about gguf_context; the helpers here are how loader -// code and family handlers talk to a gguf_context once it's open. +// INTERNAL. The public C ABI knows nothing about gguf_context; these +// helpers are how loader code and family handlers read a gguf_context once +// it's open. // -// What lives here: -// -// - KvResult: a tri-state {Absent, Ok, BadType} return type for -// low-level KV readers. The whole point is to distinguish "the -// converter didn't write this key" (Absent — usually fine for -// optional keys) from "the converter wrote it with the wrong -// type" (BadType — always a converter bug, surface as -// TRANSCRIBE_ERR_GGUF). The original 2B implementation collapsed -// these two into a single boolean and silently treated bad-type -// optional KV as missing, which papered over real converter -// mistakes. -// -// - read_*_kv low-level helpers: one per scalar / string / string- -// array KV type. Each returns KvResult and writes the parsed -// value into an out parameter on Ok. On Absent or BadType the -// out parameter is left untouched. -// -// - read_capability_kv / read_languages_kv: post-load helpers that -// family handlers call after applying their own family defaults. -// Capabilities and languages are family-agnostic in the sense -// that any family handler reads the same keys via the same -// helpers; per-family code only fills in defaults for keys the -// converter chose to omit. -// -// Why this is its own header instead of squatting on transcribe-loader.h: -// the loader is the "open a GGUF and identify it" object. These helpers -// are the "now that the GGUF is open, here is how shared metadata is -// read." Different jobs. +// - KvResult: a tri-state {Absent, Ok, BadType} that distinguishes "key +// not written" (Absent, usually fine for optional keys) from "written +// with the wrong type" (BadType, a converter bug -> TRANSCRIBE_ERR_GGUF). +// - read_*_kv low-level helpers: one per scalar / string / array KV type, +// writing the parsed value on Ok and leaving the out param untouched +// on Absent / BadType. +// - read_capability_kv / read_languages_kv: post-load helpers family +// handlers call after applying their own defaults; the schema is +// shared across families. #pragma once @@ -47,10 +28,8 @@ struct transcribe_model; namespace transcribe { -// Tri-state result for a single GGUF KV read. The strict-now intent is -// that any consumer of these helpers explicitly handles BadType — the -// type system makes "I forgot about wrong type" a compile-time -// observation rather than a silent runtime fallback. +// Tri-state result for a single GGUF KV read. Consumers must handle +// BadType explicitly rather than fall back silently. enum class KvResult { Absent, // Key not present in the gguf. Ok, // Key present, type matches expectation, value extracted. @@ -109,10 +88,7 @@ KvResult read_int32_array_kv(const gguf_context * ctx, const char * key, // --------------------------------------------------------------------------- // // These sit on top of the low-level readers and fold -// KvResult → transcribe_status, logging diagnostics with a -// caller-supplied family tag. Every per-family weights.cpp was -// carrying its own copy of these; the only difference was the log -// prefix string. +// KvResult → transcribe_status, logging with a caller-supplied family tag. // // "Required" helpers: Absent and BadType both surface as // TRANSCRIBE_ERR_GGUF. "Optional" helpers: Absent silently applies @@ -155,24 +131,13 @@ transcribe_status read_optional_int32_kv(const gguf_context * gguf, // Post-load shared metadata // --------------------------------------------------------------------------- // -// These are the "every family handler does this" pieces of post-load -// state population. They live here rather than in each family's -// arch// directory because the schema (stt.capability.*, +// Post-load shared metadata population. The schema (stt.capability.*, // general.languages) is shared across families by design. -// Read all currently-recognized stt.capability.* boolean keys into -// caps, leaving fields untouched for absent keys. Callers should call -// this AFTER their own family-default population (apply_family_invariants -// or equivalent), so KV present overrides the family default and KV -// absent leaves the family default in place. This matches PLAN.md's -// "GGUF KV is authoritative; if a key is absent, the architecture -// supplies a default" rule. -// -// In 2B the recognized keys are stt.capability.translate, -// stt.capability.lang_detect, and stt.capability.streaming. The -// timestamp granularity capability fields are deliberately deferred -// until the public-header semantics are sharpened — see RESUME.md -// open question 4. +// Read all recognized stt.capability.* boolean keys into caps, leaving +// fields untouched for absent keys. Call this AFTER family-default +// population so KV present overrides the default and KV absent keeps it. +// Recognized keys: stt.capability.{translate,lang_detect,streaming}. // // Returns: // TRANSCRIBE_OK on success or "all keys absent". diff --git a/src/transcribe-model.cpp b/src/transcribe-model.cpp index eae9f88b..909249f0 100644 --- a/src/transcribe-model.cpp +++ b/src/transcribe-model.cpp @@ -1,25 +1,13 @@ // transcribe-model.cpp - out-of-line definitions for the model and -// context base classes. -// -// Why an out-of-line .cpp at all when both bases are short: -// -// 1. Anchors the virtual destructors in one TU. Otherwise every TU -// that includes transcribe-model.h would emit a vtable + RTTI -// copy, which clang and gcc both warn about (-Wweak-vtables). -// 2. Holds set_languages(), which is non-trivial enough that inlining -// it in the header would force the / includes on -// every consumer for no benefit. +// context base classes. Anchors the virtual destructors in one TU +// (avoids -Wweak-vtables) and holds set_languages(). #include "transcribe-model.h" #include "transcribe-session.h" #include -// Anchor for the model base vtable. The default destructor is correct; -// we just need it defined in exactly one TU. transcribe_model::~transcribe_model() = default; - -// Same anchoring trick for the context base. transcribe_session::~transcribe_session() = default; void transcribe_session::clear_result() { @@ -51,11 +39,9 @@ void transcribe_session::clear_result() { } void transcribe_model::set_languages(std::vector langs) { - // Move the strings into the model so their c_str() pointers stay - // valid for the model's lifetime. Anything that previously lived in - // language_storage_ is dropped (and the corresponding entries in - // language_ptrs_ become dangling, which is why we rebuild the - // pointer vector below before exposing it through caps). + // Move the strings into the model so their c_str() pointers stay valid + // for the model's lifetime. The pointer vector is rebuilt below because + // the previous entries dangle after the storage is replaced. language_storage_ = std::move(langs); language_ptrs_.clear(); @@ -64,11 +50,8 @@ void transcribe_model::set_languages(std::vector langs) { language_ptrs_.push_back(s.c_str()); } - // Publish to the capability struct atomically from the caller's - // perspective: count and pointer match the same backing storage. - // An empty vector exposes (0, nullptr) instead of (0, &empty[0]) - // because some callers reasonably treat a non-null pointer as a - // claim that there is at least one element. + // Publish count and pointer from the same backing storage. An empty + // vector exposes (0, nullptr) rather than (0, &empty[0]). caps.n_languages = static_cast(language_storage_.size()); caps.languages = language_ptrs_.empty() ? nullptr : language_ptrs_.data(); } diff --git a/src/transcribe-model.h b/src/transcribe-model.h index b0cc4d49..1d69940c 100644 --- a/src/transcribe-model.h +++ b/src/transcribe-model.h @@ -1,28 +1,14 @@ // transcribe-model.h - internal base class for the public opaque // transcribe_model handle. // -// This header is INTERNAL. The public C ABI in include/transcribe.h only -// forward-declares `struct transcribe_model;`; the real definition lives -// here so per-family code (src/arch//...) can derive from it. -// -// Layout: hybrid of "minimal base" and "full polymorphism" — the base -// owns the universal data and per-call dispatch goes through the Arch -// trait callbacks. RESUME.md "Decisions still load-bearing" has the -// rationale and the rejected alternatives. -// -// - The base owns everything that is genuinely invariant across all -// families: the dispatch token (`arch`), the variant string, the -// bound runtime backend string, and the public capabilities struct. -// - The base has a virtual destructor so the central -// transcribe_model_free() can `delete model;` polymorphically without -// a per-family free callback. -// - Per-family subclasses (ParakeetModel, etc.) own everything else: -// the gguf_context, weights, decoder state. RAII frees them in the -// subclass destructor. -// -// The language pointer chain (the only mildly tricky part of -// transcribe_capabilities) is centralized in set_languages() so per-family -// load() code never has to get it right. +// INTERNAL. The public C ABI forward-declares `struct transcribe_model;`; +// the real definition lives here so per-family code (src/arch//...) +// can derive from it. The base owns what is invariant across all families +// (dispatch token, variant/backend strings, capabilities) and has a virtual +// destructor so transcribe_model_free() can delete polymorphically; +// per-family subclasses own the gguf_context, weights, and decoder state +// (freed via RAII in the subclass destructor). The capabilities language +// pointer chain is centralized in set_languages(). #pragma once @@ -43,9 +29,8 @@ class Tokenizer; struct ggml_backend; typedef struct ggml_backend * ggml_backend_t; -// The public C ABI forward-declares this as `struct transcribe_model;`, -// so the real definition stays in the global namespace and uses the -// `struct` keyword for ABI compatibility with C callers. +// Defined in the global namespace with the `struct` keyword to match the +// public C ABI forward declaration `struct transcribe_model;`. struct transcribe_model { // Dispatch token. Set by per-family load() to the family's Arch // instance. The central dispatcher reads this for arch_string() and @@ -66,10 +51,8 @@ struct transcribe_model { std::map meta; // Runtime backend currently bound to this model. Empty string means - // "no backend bound" — both pre-binding and the model == NULL case - // collapse to that one stable rule. See the public header for the - // full semantics. In 2B no real backend is wired, so loaders leave - // this empty. + // "no backend bound" (both pre-binding and the model == NULL case + // collapse to that one rule). See the public header for full semantics. std::string backend; // The resolved primary compute backend this model runs on (the handle @@ -142,10 +125,8 @@ struct transcribe_model { // derived from this model — load time is a model-scoped fact. int64_t t_load_us = 0; - // Default-constructs every member via its in-class initializer - // (caps is zero-filled by `caps{}`). Explicitly defaulted because - // the deleted copy/move declarations below otherwise suppress the - // implicit default constructor. + // Explicitly defaulted because the deleted copy/move declarations + // below otherwise suppress the implicit default constructor. transcribe_model() = default; virtual ~transcribe_model(); @@ -197,15 +178,13 @@ struct transcribe_model { namespace transcribe { -// Internal feature-bit helpers. Per-family load() / capability KV -// reader calls set_feature; the central probe (transcribe_model_supports) -// reads via has_feature. Both treat out-of-range values as no-ops / false -// so the bitset stays bounded and unused bits remain zero. +// Internal feature-bit helpers. Per-family load() calls set_feature; the +// central probe (transcribe_model_supports) reads via has_feature. Both +// treat out-of-range values as no-ops / false so the bitset stays bounded. // -// The feature parameter is int, NOT transcribe_feature, on purpose: the -// public entry point receives whatever int a C caller passed, and loading -// an out-of-range value through the enum type is UB in C++. Enum constants -// convert implicitly, so internal callers are unaffected. +// The feature parameter is int, NOT transcribe_feature, on purpose: a C +// caller can pass any int, and loading an out-of-range value through the +// enum type is UB in C++. Enum constants convert implicitly. inline void set_feature(transcribe_model * m, int f, bool on) { if (m == nullptr) return; const unsigned bit = static_cast(f); diff --git a/src/transcribe-session.h b/src/transcribe-session.h index 233546e4..524b9159 100644 --- a/src/transcribe-session.h +++ b/src/transcribe-session.h @@ -1,17 +1,12 @@ // transcribe-session.h - internal base class for the public opaque // transcribe_session handle. // -// Mirror of transcribe-model.h. Same Option B layout: a small base owned -// by the central dispatch, derived per-family contexts owning everything -// else. The base has a virtual destructor so transcribe_session_free() -// can `delete ctx;` polymorphically. -// -// The result storage is on the BASE because it's family-agnostic: the -// public C ABI surfaces the same hierarchy (segments / words / tokens -// / full text) regardless of which family produced it. Per-family -// run() drivers populate these vectors during decode; the public -// accessors in transcribe.cpp read them directly with no per-family -// dispatch. +// Mirrors transcribe-model.h: a small base owned by the central dispatch, +// derived per-family contexts owning everything else, with a virtual +// destructor so transcribe_session_free() can delete polymorphically. +// Result storage lives on the base because it's family-agnostic (segments +// / words / tokens / full text); per-family run() drivers populate the +// vectors and the public accessors read them directly. #pragma once @@ -70,7 +65,7 @@ struct transcribe_session { // clamped down to the model maximum by the family that honors it. // Only hard-context-cap families read this; chunked / unbounded // families ignore it. See include/transcribe.h - // (transcribe_session_params::n_ctx) and docs/input-limits.md. + // (transcribe_session_params::n_ctx). int32_t n_ctx = 0; // Cached kv_type the caller passed at init time. Hoisted to the base so @@ -87,24 +82,12 @@ struct transcribe_session { int64_t t_encode_us = 0; int64_t t_decode_us = 0; - // ----------------------------------------------------------------- // Result storage (family-agnostic; populated by per-family run()). - // - // The shape mirrors PLAN.md "Results": a flat backing array per - // level (segments / words / tokens) plus forward and backward - // pointers between levels. v1 produces a single segment per run; - // word splitting is family-specific but the per-token / per-word - // / per-segment fields below are the same for every family. - // - // result_kind tells the public accessor what timestamp granularity - // the family populated. v1 Parakeet TDT produces token-level - // timestamps from the encoder frame indices; word/segment t0/t1 - // are computed from the contained tokens. - // - // has_result is the "transcribe_run produced a real result" flag. - // When false (no run yet, or the most recent run failed before - // populating), every accessor returns its safe sentinel ("", - // 0, NAN). + // A flat backing array per level (segments / words / tokens) with + // forward/backward indices between levels. result_kind is the + // timestamp granularity the family populated; has_result is false + // until a run populates, in which case accessors return safe + // sentinels ("", 0, NAN). struct TokenEntry { int id = 0; @@ -140,29 +123,22 @@ struct transcribe_session { std::vector segments; std::string full_text; - // ----------------------------------------------------------------- - // Offline batch results (transcribe_run_batch). - // - // The scratch fields above (tokens / words / segments / full_text / - // detected_language / result_kind / has_result) remain the single - // "current result" slot that every per-family run() writes into and - // that the legacy single-shot accessors read. Batch results are a - // separate vector so a per-family run() needs no change to participate - // in the generic serial-fallback batch path: the dispatcher calls - // run() once per utterance and snapshots the scratch slot into one - // ResultSet here. A family with a real batched run_batch() hook writes - // these entries directly. + // Offline batch results (transcribe_run_batch). The scratch fields + // above are the single "current result" slot every run() writes into + // and the single-shot accessors read; batch_results is a separate + // vector so the serial-fallback path can snapshot each utterance + // without per-family changes. // // Source-of-truth rule for the public accessors: // - batch_results non-empty -> the batch accessors index it, and - // the scratch slot is restored to mirror batch_results[0] so the - // legacy single accessors stay coherent (they show utterance 0). + // the scratch slot mirrors batch_results[0] so the single + // accessors stay coherent (they show utterance 0). // - batch_results empty -> single-shot mode; the batch // accessors synthesize index 0 from the scratch slot. - // batch_results is cleared at the top of every transcribe_run and - // transcribe_run_batch (NOT by clear_result, which only wipes the - // scratch slot and streaming snapshot so a per-utterance run() inside - // the fallback loop does not erase already-accumulated entries). + // Cleared at the top of every transcribe_run / transcribe_run_batch, + // NOT by clear_result (which only wipes the scratch slot + streaming + // snapshot, so a run() inside the fallback loop does not erase + // already-accumulated entries). struct ResultSet { std::vector tokens; std::vector words; @@ -215,18 +191,12 @@ struct transcribe_session { transcribe_timestamp_kind result_kind = TRANSCRIBE_TIMESTAMPS_NONE; bool has_result = false; - // ----------------------------------------------------------------- // Abort / cancellation (set via transcribe_set_abort_callback). - // - // Per-family run() drivers call poll_abort() at chunk boundaries - // and between greedy decode steps. A callback returning true sets - // was_aborted and the run path returns TRANSCRIBE_ERR_ABORTED, - // preserving whatever partial segments accumulated before the - // abort point. - // - // was_aborted is cleared at the top of every transcribe_run; it - // is NOT cleared by clear_result because the result may be - // explicitly preserved after an abort (partial-retained contract). + // run() drivers call poll_abort() at chunk / decode-step boundaries; + // a callback returning true sets was_aborted and the run returns + // TRANSCRIBE_ERR_ABORTED with partial segments preserved. was_aborted + // is cleared at the top of every transcribe_run, NOT by clear_result + // (the partial result may be deliberately retained). transcribe_abort_callback abort_cb = nullptr; void * abort_userdata = nullptr; bool was_aborted = false; @@ -239,33 +209,21 @@ struct transcribe_session { return false; } - // Set by a per-family run() driver when greedy decode stops because it - // reached the model's context / position cap before end-of-stream, - // i.e. the output transcript was truncated. The partial result is - // retained. Surfaced via transcribe_was_truncated(); like was_aborted - // it is cleared at the top of every transcribe_run (NOT by - // clear_result, since the partial result is preserved). Distinct from - // the up-front TRANSCRIBE_ERR_INPUT_TOO_LONG rejection: this is the - // "couldn't finish", that is the "couldn't start". See - // docs/input-limits.md. + // Set by a run() driver when decode stops at the model's context / + // position cap before end-of-stream (output truncated; partial result + // retained). Surfaced via transcribe_was_truncated(); cleared at the + // top of every transcribe_run, NOT by clear_result. Distinct from the + // up-front TRANSCRIBE_ERR_INPUT_TOO_LONG rejection (couldn't finish vs + // couldn't start). See docs/input-limits.md. bool was_truncated = false; - // ----------------------------------------------------------------- - // Streaming state. - // - // Lifecycle (stream_state) is separated from result snapshot so - // that clear_result() — which is called by both transcribe_run and - // transcribe_stream_begin — can wipe per-call data without churning - // the IDLE/ACTIVE/FINISHED/FAILED state machine. The dispatcher - // manages stream_state explicitly at begin/feed/finalize/reset. - // - // The snapshot fields below (revision, committed counts, - // last_status, audio cursors) ARE cleared by clear_result so that - // a fresh transcribe_run starts with zeroed streaming bookkeeping - // and a streaming caller does not see stale counters across runs. - // - // Audio cursors are us-precision; the public stream_update struct - // exposes them as ms. + // Streaming state. Lifecycle (stream_state) is separated from the + // result snapshot so clear_result() can wipe per-call data without + // churning the IDLE/ACTIVE/FINISHED/FAILED machine, which the + // dispatcher manages explicitly. The snapshot fields below (revision, + // committed counts, last_status, audio cursors) ARE cleared by + // clear_result. Audio cursors are us-precision; the public + // stream_update struct exposes them as ms. transcribe_stream_state stream_state = TRANSCRIBE_STREAM_IDLE; int32_t stream_revision = 0; int n_committed_segments = 0; @@ -278,16 +236,13 @@ struct transcribe_session { TRANSCRIBE_STREAM_COMMIT_AUTO; uint32_t stream_stable_prefix_agreement_n = 0; - // Session-owned copies of the caller's run-params strings, refreshed on - // every transcribe_stream_begin. The dispatcher hands the family hooks a - // params view whose language/target_language point HERE, so a family - // that captures `*run_params` (parakeet re-reads .language on every - // feed; moonshine_streaming carries the copy through its decode path) - // holds pointers into library-owned storage. The public contract lets - // the caller free every params pointer the moment begin returns — - // retaining a caller pointer is a use-after-free. Stable for the - // stream's lifetime: only the next begin mutates these, and it also - // refreshes every retained copy. + // Session-owned copies of the caller's run-params strings, refreshed + // on every transcribe_stream_begin. The dispatcher hands the family + // hooks a params view whose language/target_language point HERE, so a + // family that captures *run_params holds pointers into library-owned + // storage (the public contract lets the caller free its params pointers + // the moment begin returns). Stable for the stream's lifetime; only the + // next begin mutates them. std::string stream_language_owned; std::string stream_target_language_owned; diff --git a/src/transcribe-tokenizer.cpp b/src/transcribe-tokenizer.cpp index 189a9e2d..fdc01636 100644 --- a/src/transcribe-tokenizer.cpp +++ b/src/transcribe-tokenizer.cpp @@ -485,7 +485,7 @@ transcribe_status Tokenizer::encode(const std::string & text, return TRANSCRIBE_OK; } - // Stage 1: pretokenize into byte-encoded "words". Dispatch on the + // Pretokenize into byte-encoded "words". Dispatch on the // pretokenizer flavor (set at load time from tokenizer.ggml.pre, // or overridden by a per-family loader via set_pretokenizer). // Unknown flavors fall back to qwen2 with a warning — keeps old @@ -504,7 +504,7 @@ transcribe_status Tokenizer::encode(const std::string & text, words = unicode::pretokenize_qwen2(text); } - // Stage 2: run the BPE merge loop per word and collect ids. + // Run the BPE merge loop per word and collect ids. out_ids.reserve(words.size() * 2); std::vector symbols; @@ -561,12 +561,10 @@ transcribe_status Tokenizer::encode(const std::string & text, ++index; } - // Seed initial bigrams. for (int i = 1; i < static_cast(symbols.size()); ++i) { push_bigram(i - 1, i); } - // Greedy merge loop. while (!queue.empty()) { Bigram bg = queue.top(); queue.pop(); @@ -595,7 +593,6 @@ transcribe_status Tokenizer::encode(const std::string & text, push_bigram(bg.left, left.next); } - // Walk the surviving chain and emit ids. for (int i = 0; i != -1; i = symbols[i].next) { const Symbol & s = symbols[i]; if (s.n == 0) continue; @@ -672,10 +669,10 @@ transcribe_status Tokenizer::load(const gguf_context * gguf) { decode_mode_ = DecodeMode::SentencePiece; } - // Optional: pretokenizer flavor. Absent means "qwen2" (historical - // default for the "gpt2" model tag on Qwen3-ASR GGUFs). Per-family - // loaders override after load() when the source file did not emit - // the key but the family's pretokenizer is fixed (Whisper → "gpt2"). + // Optional: pretokenizer flavor. Absent means "qwen2" (the default + // for the "gpt2" model tag on Qwen3-ASR GGUFs). Per-family loaders + // override after load() when the source file did not emit the key but + // the family's pretokenizer is fixed (Whisper → "gpt2"). pre_.clear(); switch (read_string_kv(gguf, "tokenizer.ggml.pre", pre_)) { case KvResult::Absent: @@ -745,12 +742,11 @@ transcribe_status Tokenizer::load(const gguf_context * gguf) { break; } - // Optional: merges. Only "gpt2" uses them; SentencePiece - // tokenizers ("unigram"/"bpe") encode via the score-based lattice - // we haven't wired up yet, and their decode path doesn't need - // merges. Missing merges on a "gpt2" file means encode() will - // fail at call time -- that's the right surface because the - // decode path still works. + // Optional: merges. Only "gpt2" uses them; SentencePiece tokenizers + // ("unigram"/"bpe") encode via a score-based lattice (not + // implemented) and their decode path needs no merges. Missing merges + // on a "gpt2" file means encode() fails at call time -- the right + // surface, since the decode path still works. merge_rank_.clear(); std::vector merges; switch (read_string_array_kv(gguf, "tokenizer.ggml.merges", merges)) { diff --git a/src/transcribe-tokenizer.h b/src/transcribe-tokenizer.h index f9460dd5..689c72b5 100644 --- a/src/transcribe-tokenizer.h +++ b/src/transcribe-tokenizer.h @@ -1,39 +1,25 @@ // transcribe-tokenizer.h - internal tokenizer (decode + encode). // -// This header is INTERNAL. The public C ABI does not expose tokenizer -// details directly; per-family models hold a Tokenizer instance and -// the decode driver plus the chat-template prompt builders read / -// write through it. +// INTERNAL. The public C ABI does not expose tokenizer details; per- +// family models hold a Tokenizer instance that the decode driver and +// chat-template prompt builders read/write through. // // Supported GGUF tokenizer models (tokenizer.ggml.model): // "unigram" - SentencePiece unigram (NeMo Parakeet, Cohere ASR) // "bpe" - SentencePiece BPE (same ▁ decode convention) // "gpt2" - llama.cpp's tag for Hugging Face byte-level BPE // (Qwen2/Qwen3 family). Loads merges at init and -// supports encode() for arbitrary UTF-8 text via the -// Qwen2 pretokenizer + byte-level BPE merge loop. +// supports encode() via the pretokenizer + BPE merge loop. // // What this class does: -// -// - load(): read a fixed set of tokenizer.ggml.* keys from a -// gguf_context. The set is intentionally the llama.cpp / -// whisper.cpp intersection so converters and bindings can reuse -// existing tooling. -// - id <-> piece lookup. find() is O(1) via a hash built at load. -// - decode(): join a token-id sequence into a string. SentencePiece -// flavors substitute U+2581 ("▁") with ASCII space; "gpt2" inverts -// the byte-level mapping to recover raw UTF-8 bytes. -// - encode(): UTF-8 text -> token ids. Implemented for "gpt2" only -// in 2B (other flavors return NOT_IMPLEMENTED). The encoder runs -// the Qwen2 pretokenizer, byte-level encodes each pretoken, and -// greedily applies BPE merges in rank order. -// -// The encoder only exists to render the Qwen3-ASR chat-template -// assistant prefix ("language {Name}") at decode prep time. -// If that's the *only* live consumer, a converter-side pre- -// tokenization would also work; we chose runtime encode so that -// forthcoming features (non-empty system prompt, streaming prefix -// rollback, forced-aligner text input) can reuse the same path. +// - load(): read a fixed set of tokenizer.ggml.* keys (the llama.cpp / +// whisper.cpp intersection, so converters can reuse existing tooling). +// - id <-> piece lookup; find() is O(1) via a hash built at load. +// - decode(): SentencePiece flavors substitute U+2581 ("▁") with ASCII +// space; "gpt2" inverts the byte-level mapping to recover UTF-8 bytes. +// - encode(): UTF-8 text -> token ids ("gpt2" only; others return +// NOT_IMPLEMENTED). Runs the pretokenizer, byte-level encodes each +// pretoken, then greedily applies BPE merges in rank order. #pragma once @@ -54,47 +40,28 @@ class Tokenizer { // Read tokenizer.ggml.* keys from a gguf_context. // - // Required keys: - // tokenizer.ggml.model (string) - "unigram", "bpe", or "gpt2". - // Other model strings are rejected with - // TRANSCRIBE_ERR_NOT_IMPLEMENTED. - // tokenizer.ggml.tokens (array of string) - the vocabulary. - // Must be non-empty. + // Required: + // tokenizer.ggml.model (string) - "unigram"/"bpe"/"gpt2". + // tokenizer.ggml.tokens (array string) - vocab, must be non-empty. // - // Optional keys (parsed if present, ignored otherwise): - // tokenizer.ggml.pre (string) - pretokenizer flavor used - // by encode() when model == "gpt2". - // Recognized values: "qwen2" (default - // when absent, for historical - // compatibility with Qwen3-ASR - // GGUFs) and "gpt2" (required for - // Whisper parity — the two split - // digit / contraction / whitespace - // runs differently). Per-family - // loaders MAY override the absent - // default after load() via - // set_pretokenizer(). - // tokenizer.ggml.scores (array of float32) - // tokenizer.ggml.token_type (array of int32) - // tokenizer.ggml.merges (array of string) - "left right" - // merge pairs in rank order. - // Required for encode() when - // model == "gpt2". - // tokenizer.ggml.unknown_token_id (uint32 / int32) - // tokenizer.ggml.bos_token_id (uint32 / int32) - // tokenizer.ggml.eos_token_id (uint32 / int32) - // tokenizer.ggml.blank_token_id (uint32 / int32) - the CTC / - // transducer blank id. + // Optional (parsed if present): + // tokenizer.ggml.pre (string) - pretokenizer flavor for + // encode() when model == "gpt2". + // "qwen2" (default when absent) or + // "gpt2" (Whisper parity; the two split + // digit/contraction/whitespace runs + // differently). Per-family loaders MAY + // override via set_pretokenizer(). + // tokenizer.ggml.scores (array float32) + // tokenizer.ggml.token_type (array int32) + // tokenizer.ggml.merges (array string) - "left right" pairs + // in rank order; required for "gpt2" + // encode(). + // tokenizer.ggml.{unknown,bos,eos,blank}_token_id (uint32 / int32) // - // Returns: - // TRANSCRIBE_OK on success. - // TRANSCRIBE_ERR_INVALID_ARG if gguf is null. - // TRANSCRIBE_ERR_GGUF if a required key is missing, - // has the wrong type, or arrays - // are length-inconsistent. - // TRANSCRIBE_ERR_NOT_IMPLEMENTED if the tokenizer model string - // is recognized but not yet - // supported (e.g. "wordpiece"). + // Returns: TRANSCRIBE_OK; INVALID_ARG (null gguf); GGUF (required key + // missing / wrong type / length-inconsistent arrays); NOT_IMPLEMENTED + // (model string recognized but unsupported, e.g. "wordpiece"). transcribe_status load(const gguf_context * gguf); // Optional special-token ids for the decode-only constructors. The diff --git a/src/transcribe-unicode.h b/src/transcribe-unicode.h index 1c6cace2..655261e0 100644 --- a/src/transcribe-unicode.h +++ b/src/transcribe-unicode.h @@ -58,10 +58,9 @@ struct CptFlags { // UTF-8 codec. Mirrors the narrow subset of llama.cpp's unicode.cpp // helpers that the pretokenizer needs. Invalid UTF-8 throws -// std::invalid_argument; callers should only pass text that made it -// through the GGUF tokenizer / Python reference already, so in -// practice we never see malformed input -- the throw is a -// programming-bug backstop, not a user-facing error path. +// std::invalid_argument; callers only pass text that already made it +// through the GGUF tokenizer / Python reference, so the throw guards a +// programming bug rather than a user-facing error path. std::string cpt_to_utf8 (uint32_t cpt); uint32_t cpt_from_utf8 (const std::string & utf8, size_t & offset); std::vector cpts_from_utf8(const std::string & utf8); diff --git a/src/transcribe-weights-util.cpp b/src/transcribe-weights-util.cpp index 305715b8..0fee0e77 100644 --- a/src/transcribe-weights-util.cpp +++ b/src/transcribe-weights-util.cpp @@ -1,10 +1,7 @@ // transcribe-weights-util.cpp - shared GGUF tensor validation helpers. // // Implementation of transcribe::weights::find_tensor. See the header -// for rationale. This is a verbatim hoist of the function that used -// to live in both src/arch/parakeet/weights.cpp and -// src/arch/cohere/weights.cpp (they were character-for-character -// identical apart from the "parakeet:"/"cohere:" log prefix). +// for rationale. #include "transcribe-weights-util.h" #include "transcribe-log.h" diff --git a/src/transcribe-weights-util.h b/src/transcribe-weights-util.h index e37f1b3c..57c072f8 100644 --- a/src/transcribe-weights-util.h +++ b/src/transcribe-weights-util.h @@ -79,12 +79,9 @@ ggml_tensor * find_tensor(ggml_context * ctx_meta, // Format a per-layer tensor name. The caller passes a printf fmt // containing exactly one %d; the layer index is substituted. Returns // a pointer to a thread-local static buffer — the next call within the -// same thread invalidates the previous pointer. -// -// This previously lived duplicated in every per-family weights.cpp. -// Keeping it as const char * avoids touching every GET_* macro call -// site (find_tensor takes const char *). The 128-byte buffer is ample -// for any realistic layer name (longest is ~50 chars + 3-digit index). +// same thread invalidates the previous pointer. const char * (not +// std::string) so GET_* macro call sites pass straight to find_tensor. +// The 128-byte buffer is ample for any realistic layer name. inline const char * lname(const char * fmt, int layer_idx) { thread_local char buf[128]; std::snprintf(buf, sizeof(buf), fmt, layer_idx); diff --git a/src/transcribe.cpp b/src/transcribe.cpp index f8a0cefd..47723963 100644 --- a/src/transcribe.cpp +++ b/src/transcribe.cpp @@ -2,23 +2,16 @@ // // What lives here: // - The C entry points declared in include/transcribe.h. -// - The single dispatch site that maps an opened GGUF file to its -// per-family Arch handler (transcribe_model_load_file). +// - The single dispatch site mapping an opened GGUF to its per-family +// Arch handler (transcribe_model_load_file). // - The log-sink publication / emission helpers. // - The factory functions that initialize each public params struct. // -// What does NOT live here: -// - GGUF reading itself (transcribe-loader.{h,cpp}). -// - Per-family load / context / run code (src/arch//...). -// - Anything ggml-specific (the public header is ggml-free; this file -// forwards to the loader which is the only place gguf.h is included). -// -// 2B: model and context lifecycle entry points are real now. The base -// classes in transcribe-model.h / transcribe-session.h hold the data -// the introspection accessors need (caps, variant, backend, arch -// name), so this file reads them directly without going through the -// per-family Arch trait. Per-family code still owns load(), -// init_context(), and run(). +// What does NOT live here: GGUF reading (transcribe-loader), per-family +// load / context / run code (src/arch//...), and ggml-specific +// code. The introspection accessors read the base structs in +// transcribe-model.h / transcribe-session.h directly; per-family code owns +// load(), init_context(), and run(). #include "transcribe.h" #include "transcribe/whisper.h" @@ -75,9 +68,7 @@ #include #include -// --------------------------------------------------------------------------- // Status -// --------------------------------------------------------------------------- // Parameter is `int` (not `transcribe_status`) so a caller passing an // out-of-range value does not form a bogus enum object. See the comment @@ -108,15 +99,11 @@ extern "C" const char * transcribe_status_string(int status) { } } -// --------------------------------------------------------------------------- // Version -// --------------------------------------------------------------------------- -// TRANSCRIBE_COMMIT is stamped by the build: the top-level CMakeLists.txt -// captures `git rev-parse --short HEAD` at configure time and passes it as a -// compile definition. Fall back to "unknown" for builds without git metadata -// (source tarballs, etc.) so the accessor never returns NULL. TRANSCRIBE_VERSION -// comes from , the single source of truth for the version. +// TRANSCRIBE_COMMIT is stamped by the build (git rev-parse --short HEAD at +// configure time); fall back to "unknown" so the accessor never returns +// NULL. TRANSCRIBE_VERSION comes from . #ifndef TRANSCRIBE_COMMIT #define TRANSCRIBE_COMMIT "unknown" #endif @@ -129,17 +116,13 @@ extern "C" const char * transcribe_version_commit(void) { return TRANSCRIBE_COMMIT; } -// --------------------------------------------------------------------------- // Raw enum reads at the public ABI boundary -// --------------------------------------------------------------------------- // -// C callers can store ANY int in an enum-typed ABI field or pass any int as -// an enum-typed argument — in C that is well-defined. In C++ (this TU), -// loading an out-of-range value through an enum-typed lvalue is undefined -// behavior, and UBSan traps it. So every public entry point reads -// caller-supplied enum values as raw bytes first, validates the raw int, and -// only then (if ever) converts to the enum type. All public enums are -// int-sized; the static_assert below pins one representative. +// C callers can store ANY int in an enum-typed ABI field; in C++ loading an +// out-of-range value through an enum-typed lvalue is UB (UBSan traps it). So +// every public entry point reads caller-supplied enum values as raw bytes +// first, validates the raw int, and only then converts to the enum type. All +// public enums are int-sized; the static_assert below pins one. static inline int enum_field_raw(const void * field) { int raw; @@ -149,13 +132,9 @@ static inline int enum_field_raw(const void * field) { static_assert(sizeof(transcribe_backend_request) == sizeof(int), "public enums must be int-sized for raw boundary reads"); -// --------------------------------------------------------------------------- -// ABI metadata -// --------------------------------------------------------------------------- -// -// sizeof/alignof for the public structs, so a binding can verify its own -// layout against this build before constructing instances. See the contract in -// . Keep this switch in sync with the transcribe_abi_struct enum. +// ABI metadata: sizeof/alignof for the public structs, so a binding can +// verify its layout against this build. Keep these switches in sync with the +// transcribe_abi_struct enum. extern "C" size_t transcribe_abi_struct_size(transcribe_abi_struct which) { switch (enum_field_raw(&which)) { @@ -199,18 +178,9 @@ extern "C" size_t transcribe_abi_struct_align(transcribe_abi_struct which) { namespace { -// Ordered rank for timestamp granularities, used by the dispatcher to -// compare a request against a family's advertised maximum. -// -// NONE < SEGMENT < WORD < TOKEN -// 0 1 2 3 -// -// AUTO is a sentinel at the call site and is handled by the caller -// (resolved to the model's max before ranking). A value that isn't a -// known enumerator gets rank -1 so an out-of-range request compares -// strictly less than every valid max (the dispatcher's earlier -// enum-range switch has already rejected those anyway; this keeps -// the helper total). +// Ordered rank for timestamp granularities (NONE < SEGMENT < WORD < TOKEN), +// used to compare a request against a family's advertised maximum. AUTO and +// any unknown value get rank -1. int timestamp_rank(transcribe_timestamp_kind k) { switch (k) { case TRANSCRIBE_TIMESTAMPS_NONE: return 0; @@ -333,47 +303,22 @@ transcribe_status validate_run_params_common( } // namespace -// --------------------------------------------------------------------------- // Logging -// --------------------------------------------------------------------------- -// -// Two-atomic log sink. Publication of (cb, userdata) stores userdata -// (relaxed) and then cb (release). Emission acquire-loads cb, early-outs -// on null, and relaxed-loads userdata. The acquire/release pairing on -// the cb pointer establishes a happens-before edge between the install -// thread and any later emitter, which is the property the supported -// "install once at startup" usage model needs. -// -// What this implementation does NOT promise: -// - Pair-atomic publication of (cb, userdata) under reconfiguration. -// Two concurrent calls to transcribe_log_set can interleave their -// stores in modification order such that an emitter observes a cb -// from one generation and a userdata from another. Even with a -// single writer, an emitter that still observes the previous cb -// may already see the new userdata via its relaxed load. -// - Lifetime safety for old sinks across reconfiguration. After -// transcribe_log_set returns, an in-flight emission on another -// thread may already be holding the previous (cb, userdata) on its -// stack and about to call through it. A caller that frees the old -// userdata immediately after reconfiguring would create a UAF. // -// The supported model (per the threading contract in transcribe.h) is -// "install once at startup, before threads or models exist." In that -// model neither caveat above can occur. Anything beyond startup-time -// install is unsupported in 0.x and is the caller's problem; the only -// guarantee the implementation provides in that case is the absence of -// data races, not correct semantics. +// Two-atomic log sink. Publication stores userdata (relaxed) then cb +// (release); emission acquire-loads cb and relaxed-loads userdata. The +// acquire/release pairing on cb is the happens-before edge the supported +// "install once at startup, before threads or models exist" model needs. +// Reconfiguration is unsupported (the (cb, userdata) pair can tear and an +// in-flight emitter may still hold the old sink); the only guarantee then is +// absence of data races, not correct semantics. // -// Three sink states, all classified from a single acquire load of -// g_log_cb: -// nullptr never configured -> stderr default for surfaced -// messages (emit_or_stderr) -// sentinel explicitly disabled via transcribe_log_set(NULL) -// -> drop everything -// anything else caller's callback -> route everything -// The sentinel exists so "caller asked for silence" is distinguishable -// from "caller never called us" without a second atomic whose pairing -// could tear under the (unsupported) reconfiguration case. +// Three sink states, classified from one acquire load of g_log_cb: +// nullptr never configured -> stderr default (emit_or_stderr) +// sentinel disabled via transcribe_log_set(NULL) -> drop everything +// anything else caller's callback -> route everything +// The sentinel distinguishes "asked for silence" from "never called us" +// without a second atomic whose pairing could tear. namespace { @@ -497,25 +442,12 @@ void log_msg(transcribe_log_level level, const char * fmt, ...) { } } // namespace transcribe -// --------------------------------------------------------------------------- // Params init functions -// --------------------------------------------------------------------------- // -// Every input params struct is initialized by zero-filling and stamping -// struct_size. This works because each field's documented default IS its -// zero value (TASK_TRANSCRIBE, TIMESTAMPS_NONE, PNC/ITN_MODE_DEFAULT, -// BACKEND_AUTO, KV_TYPE_AUTO are all 0; gpu_device 0 = auto; NULL -// pointers; keep_special_tags false = strip). That invariant lets the -// init functions stay as memset + stamp without per-field assignments. -// `{0}` itself is NOT accepted as a defaults form — struct_size == 0 is -// rejected — because uninitialized stack memory can silently hit the -// zero case; callers reach defaults via NULL only. -// -// These are one-argument by design: they assume the caller and library -// agree on the struct layout, which holds for the supported build-with- -// your-consumers distribution model. A caller passing a struct smaller -// than the library's view would be a version-skew bug handled at release -// time (soname), not in this API. +// Each input params struct is initialized by zero-filling and stamping +// struct_size, which works because every field's documented default is its +// zero value. struct_size == 0 is NOT a defaults form (uninitialized stack +// memory could hit the zero case); callers reach defaults via NULL only. extern "C" void transcribe_model_load_params_init(struct transcribe_model_load_params * p) { if (p == nullptr) { return; } @@ -545,18 +477,13 @@ extern "C" void transcribe_stream_params_init(struct transcribe_stream_params * p->struct_size = sizeof(*p); } -// --------------------------------------------------------------------------- // Output struct init functions -// --------------------------------------------------------------------------- // // Output structs are caller-allocated buffers the library writes into, -// bounded by min(caller struct_size, library size). The caller must -// declare its buffer size via struct_size; these init functions do that -// and zero-fill the rest. Every output field is designed so its zero -// value means "absent / unknown / false / none", so a zeroed struct + -// struct_size is the correct empty state. (Unlike input params, a {0} -// output struct is NOT accepted by the accessors — struct_size == 0 there -// is a real "you forgot to init the buffer" error.) +// bounded by min(caller struct_size, library size). Every output field's +// zero value means "absent / unknown / false / none", so a zeroed struct + +// struct_size is the correct empty state. struct_size == 0 is rejected by +// the accessors as a "you forgot to init the buffer" error. extern "C" void transcribe_capabilities_init(struct transcribe_capabilities * p) { if (p == nullptr) { return; } @@ -612,21 +539,13 @@ extern "C" void transcribe_token_init(struct transcribe_token * p) { p->struct_size = sizeof(*p); } -// Whisper telemetry + run-extension init and chunk-trace accessors live -// in arch/whisper/public.cpp (family source) so this dispatcher stays -// family-agnostic, matching parakeet/moonshine. - -// The whisper/sensevoice/funasr_nano/canary per-family run-params -// factories were removed in the Phase-2 migration: -// - whisper's knobs are now in transcribe_whisper_run_ext (reached via -// transcribe_run_params::family); use transcribe_whisper_run_ext_init(). -// - sensevoice/funasr_nano `use_itn` collapsed into the generic -// transcribe_run_params::itn enum. -// - canary `pnc` collapsed into the generic transcribe_run_params::pnc enum. +// Whisper telemetry + run-extension init and chunk-trace accessors live in +// arch/whisper/public.cpp so this dispatcher stays family-agnostic. Whisper +// run knobs are in transcribe_whisper_run_ext (via run_params::family); +// sensevoice/funasr_nano use_itn and canary pnc use the generic run_params +// itn/pnc enums. -// --------------------------------------------------------------------------- // Extension helpers -// --------------------------------------------------------------------------- extern "C" transcribe_status transcribe_ext_check( const struct transcribe_ext * ext, @@ -749,21 +668,18 @@ static bool valid_stream_commit_policy(int policy) { } // namespace -// --------------------------------------------------------------------------- // Backend modules and device discovery -// --------------------------------------------------------------------------- // // transcribe_init_backends loads ggml backend modules from ONE caller-named -// directory (a Python wheel's package dir, an app bundle, ...). It never -// widens the search: ggml_backend_load_all_from_path with a non-null path -// scans only that path. ggml tolerates every per-module failure (missing -// system deps such as the Vulkan loader -> dlopen fails -> module skipped), -// which is exactly the ship-Vulkan-by-default degradation contract. +// directory (a Python wheel's package dir, an app bundle, ...) via +// ggml_backend_load_all_from_path; it never widens the search. ggml tolerates +// every per-module failure (missing system deps -> dlopen fails -> module +// skipped), the ship-Vulkan-by-default degradation contract. // // NOTE: these definitions must sit at FILE scope, outside the anonymous // namespace above. clang exports extern "C" functions defined inside an -// unnamed namespace; gcc gives them internal linkage, which strips them -// from the shared library on Linux (undefined references at link). +// unnamed namespace; gcc gives them internal linkage, which strips them from +// the shared library on Linux (undefined references at link). // ggml's MSVC timer (ggml_time_us/ms in ggml.c) divides QueryPerformanceCounter // by a global frequency that ggml_time_init() fills in from @@ -873,13 +789,10 @@ extern "C" transcribe_status transcribe_init_backends(const char * artifact_dir) ggml_backend_load_all_from_path(artifact_dir); s_loaded_dirs.insert(key); - // Post-scan device summary, through the installed sink only. Release - // ggml compiles its per-module load-failure diagnostics out - // (silent=true under NDEBUG), so this one line is the reliable - // signal a host's log callback gets for "which backends made it" — - // the input to debugging an accelerator that degraded to CPU. Sent - // via the callback-only emitter on purpose: a host without a sink - // keeps today's stderr behavior, no new import-time noise. + // Post-scan device summary, through the installed sink only (release + // ggml compiles its per-module load-failure diagnostics out, so this + // line is the reliable signal of which backends made it). Callback-only + // so a host without a sink gets no new import-time stderr noise. char devices[256]; size_t off = 0; const size_t n_dev = ggml_backend_dev_count(); @@ -1367,19 +1280,12 @@ static void publish_observable_delta( } } -// Advisory pnc/itn warning. Emits a WARN log message when a non-DEFAULT -// caller request hits a model that does not support runtime control of -// the corresponding axis, then returns so the dispatcher can proceed. -// "Best effort" semantics: the model produces *something* either way, -// and silently ignoring the request would be a footgun. The reserved -// TRANSCRIBE_ERR_UNSUPPORTED_PNC / _ITN codes are NOT returned today; -// they are placeholders for a future opt-in strict-advisory mode (a -// per-call advisory_strict flag on transcribe_run_params). -// -// The arch + variant strings are included in the message so a grepping -// developer can pinpoint which model dropped the request without having -// to reproduce the call. Emission falls back to stderr when no log -// callback is installed (mirrors transcribe_print_timings). +// Advisory pnc/itn warning. Emits a WARN when a non-DEFAULT request hits a +// model that does not support runtime control of that axis, then returns so +// the dispatcher proceeds with best-effort semantics. The reserved +// TRANSCRIBE_ERR_UNSUPPORTED_PNC / _ITN codes are NOT returned today +// (placeholders for a future opt-in strict mode). The message includes the +// arch + variant strings to pinpoint which model dropped the request. void warn_unsupported_advisory( const struct transcribe_model * model, const struct transcribe_run_params * rp) @@ -1426,9 +1332,7 @@ void warn_unsupported_advisory( } // namespace -// --------------------------------------------------------------------------- -// Lifecycle (stubs) -// --------------------------------------------------------------------------- +// Lifecycle extern "C" transcribe_status transcribe_model_load_file( const char * path, @@ -1484,17 +1388,10 @@ extern "C" transcribe_status transcribe_model_load_file( return TRANSCRIBE_ERR_INVALID_ARG; } - // Stage 0: format sniff. Two formats are supported today: - // GGUF magic 0x46554747 ("GGUF") — the canonical path. - // ggml magic 0x67676d6c — legacy whisper.cpp `.bin`. - // We detect by reading the first four bytes rather than by file - // extension: HF distributes whisper.cpp models under the `.bin` - // suffix even though some are GGUFs internally, and conversely a - // user could rename a .bin to .gguf. - // - // Open semantics preserve the existing public contract: missing - // path → ERR_FILE_NOT_FOUND; everything else (truncated header, - // unrecognized magic, structural failure) collapses to ERR_GGUF. + // Format sniff by reading the first four bytes (not the file extension, + // which HF/users mislabel): GGUF magic 0x46554747 is the canonical path, + // ggml magic 0x67676d6c is a legacy whisper.cpp .bin. Public contract: + // missing path -> ERR_FILE_NOT_FOUND; every other failure -> ERR_GGUF. { struct stat st {}; if (::stat(path, &st) != 0) { @@ -1523,29 +1420,23 @@ extern "C" transcribe_status transcribe_model_load_file( } } - // 0x67676d6c ("ggml" little-endian): legacy whisper.cpp .bin. - // Hand off to the whisper-specific .bin adapter, which validates - // the hparams as whisper-shaped (rejecting Silero VAD and other - // unrelated `ggml`-magic files with UNSUPPORTED_ARCH). + // 0x67676d6c ("ggml" little-endian): legacy whisper.cpp .bin. Hand off + // to the whisper .bin adapter, which validates the hparams as + // whisper-shaped (rejecting unrelated ggml-magic files like Silero VAD). if (magic == 0x67676d6cu) { return transcribe::whisper::load_from_bin(path, params, out_model); } - // Stage 1: header-only GGUF inspection. The loader returns - // FILE_NOT_FOUND for a missing path (pre-checked via stat) or - // ERR_GGUF for any structural failure inside gguf_init_from_file - // (corrupt magic, truncated header, missing general.architecture). - // The Loader is stack-allocated; if anything below this point bails - // out before transferring ownership, its destructor frees the - // gguf_context on unwind. + // Header-only GGUF inspection. The Loader is stack-allocated; if + // anything below bails before transferring ownership, its destructor + // frees the gguf_context on unwind. transcribe::Loader loader; if (const transcribe_status st = loader.open(path); st != TRANSCRIBE_OK) { return st; } - // Stage 2: per-family dispatch. The architecture string came from - // the GGUF KV section so it is guaranteed non-null and NUL-terminated - // by the loader. + // Per-family dispatch. The architecture string came from the GGUF KV so + // the loader guarantees it is non-null and NUL-terminated. const transcribe::Arch * arch = transcribe::find_arch(loader.arch().c_str()); if (arch == nullptr) { return TRANSCRIBE_ERR_UNSUPPORTED_ARCH; @@ -1661,17 +1552,12 @@ extern "C" void transcribe_session_free(struct transcribe_session * session) { transcribe_model_free(owned); // NULL is a no-op } -// --------------------------------------------------------------------------- // Convenience: open / close / get_model -// --------------------------------------------------------------------------- // -// transcribe_open bundles load + session_init for the common single-session -// case and hands back a session that owns its model (owns_model = true). -// transcribe_close is a thin alias for transcribe_session_free — both honor -// owns_model and free the owned model after the session, so calling either -// on an open()-created session does the right thing. Pre-1.0 the alias is -// kept for source-compatibility; consumers should migrate to -// transcribe_session_free. +// transcribe_open bundles load + session_init and hands back a session that +// owns its model (owns_model = true). transcribe_close is a thin alias for +// transcribe_session_free; both honor owns_model and free the owned model +// after the session. extern "C" transcribe_status transcribe_open( const char * path, @@ -1753,24 +1639,14 @@ extern "C" bool transcribe_was_truncated(const struct transcribe_session * sessi return session->was_truncated; } -// --------------------------------------------------------------------------- // Streaming dispatcher -// --------------------------------------------------------------------------- -// -// State transitions are managed entirely here; per-family hooks see -// stream_state == ACTIVE on entry to begin/feed/finalize and never -// observe transitions themselves. Hooks own the per-utterance result -// data on the context and may freely mutate tokens/words/segments, -// low-level candidate committed counts, audio cursors, and -// stream_revision; the dispatcher owns lifecycle state, last_status, and -// the public committed/tentative text view. // -// The result snapshot and lifecycle state are deliberately separate. -// clear_result() (called by both begin and transcribe_run) wipes the -// snapshot but never touches stream_state, so a future caller that -// wants to "rewind to a fresh result without restarting the stream" -// has a path. The streaming dispatcher is currently the only caller -// that drives stream_state. +// State transitions are managed entirely here; hooks see ACTIVE on entry to +// begin/feed/finalize and never observe transitions. Hooks own the +// per-utterance result data (tokens/words/segments, committed counts, audio +// cursors, stream_revision); the dispatcher owns lifecycle state, +// last_status, and the public committed/tentative text view. clear_result() +// wipes the snapshot but never touches stream_state. extern "C" transcribe_status transcribe_stream_begin( struct transcribe_session * session, @@ -2228,24 +2104,15 @@ static transcribe_status run_one_inner( // Pre-clear family-extension SHAPE validation: a run ext with a bad // transcribe_ext header size or a kind the model does not accept must - // NOT wipe the previous result snapshot. (The family's run_validate - // hook, run below after the run-param checks, additionally enforces - // the per-kind minimum struct size.) Mirrors the transcribe_stream_begin - // contract for stream_params->family. The kind-accept probe needs a - // valid model; when model is null we skip the pre-clear check and let - // the post-clear NOT_IMPLEMENTED path handle it (the existing - // snapshot-wipe-on-NOT_IMPLEMENTED contract is preserved). + // NOT wipe the previous result snapshot (mirrors the + // transcribe_stream_begin contract). When model is null we skip the + // pre-clear check and let the post-clear NOT_IMPLEMENTED path handle it. // - // NOTE the limit of this guarantee: it covers ext *shape* (size/kind) - // and the run-param checks below. A family is free to defer deeper - // *value* validation to its run() handler, in which case a - // correctly-shaped but semantically-malformed ext can still be - // rejected after the snapshot is cleared. Whisper does exactly this - // for its prompt-semantics checks — an intentional, accepted gap on - // the one-shot run() path (see whisper_run_validate and - // docs/follow-ups.md). A family wanting full pre-clear safety should - // validate values in run_validate, the way parakeet's stream_validate - // vets its (L,C,R) menu. + // This guarantee covers ext shape (size/kind) and the run-param checks + // below. A family that defers deeper value validation to run() (whisper + // does this for prompt semantics) can still reject a correctly-shaped + // ext after the snapshot is cleared; full pre-clear safety requires + // validating values in run_validate. if (params->family != nullptr && session->model != nullptr && session->model->arch != nullptr) { @@ -2369,9 +2236,7 @@ extern "C" transcribe_status transcribe_run( return run_one_inner(session, pcm, n_samples, params); } -// --------------------------------------------------------------------------- // Batch run (offline) -// --------------------------------------------------------------------------- // // transcribe_run_batch validates the shared run_params ONCE, then either // delegates to the family's batched run_batch() fast path or falls back to @@ -2397,13 +2262,10 @@ transcribe_session::ResultSet snapshot_scratch_result( } // Pad batch_results out to `n` entries with explicit aborted failures. -// Called only on an aborted batch: retained results keep their real ResultSet; -// synthesized slots carry TRANSCRIBE_ERR_ABORTED, meaning "did not complete -// because the batch was aborted" (NOT that the utterance itself reached an -// abort checkpoint). This is what gives the result-set view one slot per input -// utterance — transcribe_batch_n_results() == n — after an ABORTED return, -// keeping the index->input mapping uniform across families and across the fast -// path vs the serial fallback. +// Called only on an aborted batch: synthesized slots carry +// TRANSCRIBE_ERR_ABORTED ("did not complete because the batch was aborted"), +// so the result-set view has one slot per input utterance +// (transcribe_batch_n_results() == n) regardless of family or path. void pad_batch_results_aborted(transcribe_session * s, int n) { while (s->batch_results.size() < static_cast(n)) { transcribe_session::ResultSet rs; @@ -2412,8 +2274,8 @@ void pad_batch_results_aborted(transcribe_session * s, int n) { } } -// Restore the scratch slot from a ResultSet so the legacy single-result -// accessors stay coherent after a batch run (they reflect utterance 0). +// Restore the scratch slot from a ResultSet so the single-result accessors +// stay coherent after a batch run (they reflect utterance 0). void restore_scratch_from_result( transcribe_session * s, const transcribe_session::ResultSet & rs) { @@ -2521,7 +2383,7 @@ extern "C" transcribe_status transcribe_run_batch( if (st == TRANSCRIBE_ERR_ABORTED) { pad_batch_results_aborted(session, n); } - // Keep the legacy single accessors coherent with utterance 0. + // Keep the single accessors coherent with utterance 0. if (!session->batch_results.empty()) { restore_scratch_from_result(session, session->batch_results.front()); } @@ -2567,22 +2429,18 @@ extern "C" transcribe_status transcribe_run_batch( pad_batch_results_aborted(session, n); } - // Restore the scratch slot to mirror utterance 0 for the legacy - // single-result accessors. clear_result() if the batch is somehow empty - // (n >= 1 guarantees at least one entry, but be defensive). + // Restore the scratch slot to mirror utterance 0 for the single-result + // accessors (n >= 1 guarantees at least one entry, but be defensive). if (!session->batch_results.empty()) { restore_scratch_from_result(session, session->batch_results.front()); } return batch_status; } -// --------------------------------------------------------------------------- // Model introspection -// --------------------------------------------------------------------------- // -// All four accessors read directly from the base struct. There is no -// per-family callback for them: per-family load() is responsible for -// filling caps / variant / backend in the base before returning success. +// These accessors read directly from the base struct; per-family load() +// fills caps / variant / backend before returning success. extern "C" transcribe_status transcribe_model_get_capabilities( const struct transcribe_model * model, @@ -2752,10 +2610,8 @@ extern "C" int transcribe_tokenize( } extern "C" const char * transcribe_model_backend(const struct transcribe_model * model) { - // Empty string means "no runtime backend bound" — see the public - // header for the full semantic. In 2B no model is ever bound to a - // backend, so per-family load() leaves this empty and we just - // pass it through. + // Empty string means "no runtime backend bound" — see the public header + // for the full semantic. if (model == nullptr) { return ""; } @@ -2775,7 +2631,7 @@ extern "C" transcribe_status transcribe_model_get_device( return st; } // The model's primary backend is bound by per-family load(); a model - // that never resolved one (or a 2B build with no real backend) has none. + // that never resolved one has none. if (model->primary_backend == nullptr) { return TRANSCRIBE_ERR_BACKEND; } @@ -2787,15 +2643,11 @@ extern "C" transcribe_status transcribe_model_get_device( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // Timings -// --------------------------------------------------------------------------- // -// All accumulators live on the base classes (transcribe_model::t_load_us -// and transcribe_session::t_{mel,encode,decode}_us). Per-family run() -// drivers populate the context-side fields; per-family load() drivers -// populate t_load_us. The accessors here just convert microseconds to -// milliseconds and return a small value type. +// Accumulators live on the base classes (transcribe_model::t_load_us and +// transcribe_session::t_{mel,encode,decode}_us), populated by per-family +// load() / run(). These accessors convert microseconds to milliseconds. extern "C" transcribe_status transcribe_get_timings( const struct transcribe_session * session, @@ -2855,16 +2707,12 @@ transcribe_reset_timings(struct transcribe_session * session) session->t_decode_us = 0; } -// --------------------------------------------------------------------------- // Result accessors -// --------------------------------------------------------------------------- // -// All accessors read from the base context's result storage -// (transcribe_session::tokens / words / segments / full_text / -// result_kind / has_result), which is populated by the per-family -// run() driver during decode. Out-of-range indices and pre-run -// access return safe sentinels per the public contract; we never -// reach into invalid memory. +// These read from the base context's result storage (tokens / words / +// segments / full_text / result_kind / has_result), populated by the +// per-family run() driver. Out-of-range indices and pre-run access return +// safe sentinels per the public contract. extern "C" const char * transcribe_full_text(const struct transcribe_session * session) { if (session == nullptr || !session->has_result) { @@ -2901,24 +2749,18 @@ extern "C" int transcribe_n_tokens(const struct transcribe_session * session) { return static_cast(session->tokens.size()); } -// --------------------------------------------------------------------------- // Result accessors - per-item rows -// --------------------------------------------------------------------------- // -// 3 copy-out accessors backed by the context's segments / words / -// tokens vectors. Each takes a caller-owned struct initialized via -// TRANSCRIBE_*_INIT and writes only the prefix that fits within the -// caller's struct_size. Out-of-range index or pre-run access leaves -// the caller's struct as zero-init (text==NULL, scalars 0); the -// dispatch path always returns OK in those cases so callers can use -// the row's text!=NULL check as the "row present" signal without -// branching on status. +// Copy-out accessors backed by the context's segments / words / tokens +// vectors, writing only the prefix that fits the caller's struct_size. +// Out-of-range index or pre-run access leaves the caller's struct zero-init +// and still returns OK, so callers can use text!=NULL as the "row present" +// signal without branching on status. // -// `text` pointers in the staged structs alias the context's -// std::string storage. The library treats that as session-owned: valid -// until the next transcribe_run / transcribe_stream_begin / -// transcribe_stream_reset / transcribe_session_free on the same -// context. The public header documents the lifetime contract. +// `text` pointers in the staged structs alias the context's std::string +// storage; the library treats that as session-owned, valid until the next +// transcribe_run / transcribe_stream_begin / transcribe_stream_reset / +// transcribe_session_free on the same context (see the public header). extern "C" transcribe_status transcribe_get_segment( const struct transcribe_session * session, @@ -3027,15 +2869,12 @@ extern "C" transcribe_status transcribe_get_token( return TRANSCRIBE_OK; } -// --------------------------------------------------------------------------- // Batch result accessors -// --------------------------------------------------------------------------- // // These index session->batch_results when a batch run populated it, and -// otherwise synthesize utterance 0 from the scratch slot after transcribe_run. -// The copy-out row accessors share the same staging shape as the single-result -// accessors above; only the source vector differs (batch ResultSet rows vs the -// scratch slot). +// otherwise synthesize utterance 0 from the scratch slot after +// transcribe_run. They share the staging shape of the single-result +// accessors above; only the source vector differs. namespace { diff --git a/tests/debug_dump_unit.cpp b/tests/debug_dump_unit.cpp index facc38e3..52492dc9 100644 --- a/tests/debug_dump_unit.cpp +++ b/tests/debug_dump_unit.cpp @@ -1,5 +1,5 @@ -// debug_dump_unit.cpp - round-trip unit test for the per-stage tensor -// dumper in src/transcribe-debug.{h,cpp}. +// debug_dump_unit.cpp - round-trip unit test for the tensor dumper in +// src/transcribe-debug.{h,cpp}. // // What we cover: // @@ -17,12 +17,12 @@ // // - Metal backend correctness. The same code path runs on Metal in // real builds (the loader binds Metal first on Apple Silicon and -// the encoder graph in step 3 will dump from Metal-resident -// activations). Phase 4 step 3 will exercise that end-to-end via -// the encoder graph; this unit test stays CPU-only so it's +// the encoder graph can dump from Metal-resident activations). +// Real-model validation exercises that end-to-end via the encoder +// graph; this unit test stays CPU-only so it's // fast and runs in CI without Metal hardware. -// - The compare_tensors.py side. That's exercised ad-hoc by the -// numerical bringup loop in step 3. +// - The compare_tensors.py side. That's exercised by numerical +// validation runs. #include "transcribe-debug.h" diff --git a/tests/decoder_smoke.cpp b/tests/decoder_smoke.cpp index 55672ec0..407a2aa4 100644 --- a/tests/decoder_smoke.cpp +++ b/tests/decoder_smoke.cpp @@ -1,9 +1,8 @@ // decoder_smoke.cpp - real-model gated end-to-end decoder accuracy test. // -// Phase 5 + phase 6 acceptance gate. Loads a real Parakeet v2 GGUF, -// runs the full pipeline (load → mel → encoder → predictor + joint -// + TDT decode → result accessor population) on samples/jfk.wav, and -// asserts: +// Loads a real Parakeet v2 GGUF, runs the full pipeline (load → mel → +// encoder → predictor + joint + TDT decode → result accessor population) +// on samples/jfk.wav, and asserts: // // 1. transcribe_run completes OK on the canonical sample. // 2. transcribe_full_text matches the canonical JFK reference text diff --git a/tests/moonshine_streaming_stream_parity.cpp b/tests/moonshine_streaming_stream_parity.cpp index 0794b236..d793542f 100644 --- a/tests/moonshine_streaming_stream_parity.cpp +++ b/tests/moonshine_streaming_stream_parity.cpp @@ -281,8 +281,7 @@ int main(int /*argc*/, char ** /*argv*/) { // Streaming parity across a range of chunk sizes. 1 ms is a // pathological "every PCM sample basically gets its own feed" // case; the others cover typical end-user buffer sizes. Throttle - // = 0 means decode on every advance (matches Phase 4b-full v0 - // behavior). + // = 0 means decode on every advance. const int chunk_ms_choices[] = { 1, 20, 40, 80, 160, 500, 1000 }; for (int chunk_ms : chunk_ms_choices) { const int chunk_samples = std::max(1, chunk_ms * 16000 / 1000); diff --git a/tests/parakeet_real_smoke.cpp b/tests/parakeet_real_smoke.cpp index a81649d0..3a9e6303 100644 --- a/tests/parakeet_real_smoke.cpp +++ b/tests/parakeet_real_smoke.cpp @@ -152,7 +152,7 @@ int main() { return EXIT_FAILURE; } - // Public ABI sanity. After phase 4 step 1, backend is one of + // Public ABI sanity. After load, backend is one of // "metal" (Apple Silicon, default) or "cpu" (fallback). The // exact label depends on the build platform; we just assert // it's non-empty and one of the expected values. diff --git a/tests/parakeet_smoke.cpp b/tests/parakeet_smoke.cpp index 0406f349..2dc922c0 100644 --- a/tests/parakeet_smoke.cpp +++ b/tests/parakeet_smoke.cpp @@ -18,7 +18,7 @@ // the per-block shape formulas baked into weights.cpp's GET() // sequence — paranoia in case the canonical shape formulas drift // out of sync with the fixture. -// 5. transcribe_model_backend() is non-empty after phase 4 step 1 +// 5. transcribe_model_backend() is non-empty after load // (Metal on Apple Silicon, CPU elsewhere or on Metal init // failure). The synthetic fixture uses the same load path as // the real model so this is the same code being exercised. @@ -133,7 +133,7 @@ int main() { return EXIT_FAILURE; } - // Public ABI sanity. After phase 4 step 1, backend is one of + // Public ABI sanity. After load, backend is one of // "metal" or "cpu" (Metal-first on Apple Silicon, CPU fallback). // The exact label is platform-dependent so we just assert // non-empty here; the real-model test does the same. diff --git a/tests/qwen3_asr_smoke.cpp b/tests/qwen3_asr_smoke.cpp index 5b729e4a..0b150492 100644 --- a/tests/qwen3_asr_smoke.cpp +++ b/tests/qwen3_asr_smoke.cpp @@ -20,7 +20,7 @@ // (see _f32_seq in make_gguf_fixtures.py). // 4. ggml_tensor::ne for representative tensors matches the // catalog shape formulas in build_qwen3_asr_weights. -// 5. The Phase 1.6 resolve_chat_tokens() call succeeded at load — +// 5. resolve_chat_tokens() succeeded at load — // every chat-template piece ended up in the resolved // ChatTokens struct with a valid id. // 6. The ffn_gate_up packed weight was populated at load (model.cpp @@ -170,7 +170,7 @@ int main() { // Capabilities: advertised language list + detection bool round- // trip; translate stays false (family invariant); max_timestamp - // stays NONE (first port doesn't emit alignment). + // stays NONE (alignment is not exposed for this family). { transcribe_capabilities caps_buf; transcribe_capabilities_init(&caps_buf); const bool caps_ok = @@ -242,7 +242,7 @@ int main() { // ----- Chat template string round-tripped ----- CHECK(!qm->chat_template.empty()); - // ----- Phase 1.6: chat-template token ids resolved at load ----- + // ----- Chat-template token ids resolved at load ----- { const auto & ct = qm->chat_tokens; CHECK_EQ_INT(ct.im_start, 2); diff --git a/tests/stream_capability_unit.cpp b/tests/stream_capability_unit.cpp index 5e99fe05..c7657ae0 100644 --- a/tests/stream_capability_unit.cpp +++ b/tests/stream_capability_unit.cpp @@ -1,5 +1,4 @@ -// stream_capability_unit.cpp - Phase 3: capability gating against a -// real loaded model. +// stream_capability_unit.cpp - capability gating against a real loaded model. // // The model-less context tests in stream_dispatch_unit.cpp exercise // the streaming dispatcher's state machine and accessor sentinels. diff --git a/tests/stream_dispatch_unit.cpp b/tests/stream_dispatch_unit.cpp index 29ee0bb8..f3c9e604 100644 --- a/tests/stream_dispatch_unit.cpp +++ b/tests/stream_dispatch_unit.cpp @@ -8,7 +8,7 @@ // // Capability-gated paths (NOT_IMPLEMENTED for non-streaming families, // TRANSLATE rejection, language validation) require a real model and -// land in Phase 3 against an existing non-streaming arch. +// are covered by stream_capability_unit.cpp. #include "transcribe-session.h" #include "transcribe-arch.h" diff --git a/tests/tokenizer_smoke.cpp b/tests/tokenizer_smoke.cpp index 998e9ed7..b1ea57c5 100644 --- a/tests/tokenizer_smoke.cpp +++ b/tests/tokenizer_smoke.cpp @@ -109,7 +109,7 @@ struct transcribe_model * load_or_fail(const char * fixture_name, } CHECK_STR_EQ(transcribe_model_arch_string(model), "parakeet"); CHECK_STR_EQ(transcribe_model_variant_string(model), expected_variant); - // After phase 4 step 1 the loader binds a runtime backend; the + // After load the loader binds a runtime backend; the // exact label is platform-dependent (Metal on Apple Silicon, CPU // elsewhere or on Metal init failure). Just assert non-empty. { diff --git a/tests/whisper_e2e_smoke.cpp b/tests/whisper_e2e_smoke.cpp index 51e98715..16a157e9 100644 --- a/tests/whisper_e2e_smoke.cpp +++ b/tests/whisper_e2e_smoke.cpp @@ -1,6 +1,6 @@ // whisper_e2e_smoke.cpp - real-model gated public-ABI test for Whisper. // -// Covers the runtime capability contract that Stage 4 must not fake: +// Covers the runtime capability contract: // language detection is advertised and exercised by running without a // language hint, default params use no-timestamp decode, explicit // SEGMENT timestamps are advertised and returned, and no-timestamp @@ -105,11 +105,8 @@ int main() { CHECK(caps->n_languages > 0); CHECK(caps->languages != nullptr); - // Stage 2 + 3 feature surface. Cancellation + temperature - // fallback + long-form shipped in Stage 2; initial_prompt - // (initial_prompt / prompt_tokens / condition_on_prev_tokens) - // shipped in Stage 3. These advisories live behind the - // transcribe_model_supports() probe, not the caps struct. + // Feature advisories live behind the transcribe_model_supports() + // probe, not the caps struct. CHECK(transcribe_model_supports(model, TRANSCRIBE_FEATURE_CANCELLATION)); CHECK(transcribe_model_supports(model, TRANSCRIBE_FEATURE_TEMPERATURE_FALLBACK)); CHECK(transcribe_model_supports(model, TRANSCRIBE_FEATURE_LONG_FORM)); @@ -192,10 +189,8 @@ int main() { // KV loop. The callback returns true after `trigger_after` // invocations; we expect TRANSCRIBE_ERR_ABORTED, transcribe_was_aborted // true, and the context to expose whatever partial content was - // committed before the abort point (Stage 1: no chunks, so - // partial == "nothing beyond the prompt-pass argmax that hasn't - // been consumed yet"). The assertion here is semantic: the run - // short-circuited and was_aborted is true. + // committed before the abort point. The assertion here is semantic: + // the run short-circuited and was_aborted is true. struct AbortState { int calls_so_far; int trigger_after; @@ -246,7 +241,7 @@ int main() { CHECK(!transcribe_was_aborted(ctx)); } - // Stage 2.6: trace accessors populated on short-form. Exactly one + // Trace accessors populated on short-form. Exactly one // chunk for a ≤ 30 s clip. Tier 0 at default params accepts, so // temperature_used == 0.0, n_fallbacks == 0. compression_ratio and // avg_logprob are real numbers from the decoder, not zero. @@ -272,7 +267,7 @@ int main() { oor.temperature_used == 0.0f && oor.n_fallbacks == 0); } - // Stage 2.2 + 2.3: long-form chunk loop with dynamic stride. + // Long-form chunk loop with dynamic stride. // Synthesize a 40-second clip from the JFK PCM (repeat pattern). // The real audio exceeds one 30-second window, so the seek loop // runs at least twice and we expect at least two chunk traces. @@ -311,7 +306,7 @@ int main() { } } - // Stage 2.5 silence-detection test is deferred: all-zero PCM is a + // The default-threshold silence-detection test is deferred: all-zero PCM is a // pathological input for the mel frontend (log10(0) behavior) and // the downstream no_speech_prob is not guaranteed to exceed the // 0.6 default threshold without real ambient audio. The mechanism @@ -319,7 +314,7 @@ int main() { // reads no_speech_triggered on every chunk. The forced-threshold // test below exercises the skip logic on JFK audio. // - // TODO (Stage 2 follow-ups; tracked on the whisper-parity branch): + // TODO: // - Real silence fixture under samples/ (room-tone or mic-off // recording) so the default-threshold no-speech path exercises // un-tuned behavior end-to-end. @@ -450,7 +445,7 @@ int main() { CHECK(!text_a.empty()); } - // ============ Stage 3: prompt + condition_on_prev_tokens ============ + // ============ Prompt + condition_on_prev_tokens ============ // ALL_SEGMENTS without condition_on_prev_tokens must reject with // INVALID_ARG before any compute runs (HF generation_whisper.py diff --git a/tests/whisper_tokenize_parity.cpp b/tests/whisper_tokenize_parity.cpp index df967dd2..5e05eb3f 100644 --- a/tests/whisper_tokenize_parity.cpp +++ b/tests/whisper_tokenize_parity.cpp @@ -1,7 +1,7 @@ // whisper_tokenize_parity.cpp - real-model gated test that // transcribe_tokenize() produces HF-identical token sequences on // Whisper's GPT-2 byte-level BPE, including the digit / contraction -// cases that forced the GPT-2 pretokenizer branch in Stage 1. +// cases that exercise the GPT-2 pretokenizer branch. // // The expected token ids were computed with HuggingFace's // openai/whisper-tiny tokenizer, add_special_tokens=False. That diff --git a/tools/transcribe-bench/main.cpp b/tools/transcribe-bench/main.cpp index c3af8f66..bb007b1c 100644 --- a/tools/transcribe-bench/main.cpp +++ b/tools/transcribe-bench/main.cpp @@ -6,14 +6,9 @@ // models and samples. Progress lines go to stderr; the JSON goes // to --json-out (or stdout if omitted). // -// Schema version: "transcribe-bench-v2". -// -// v2 changes from v1: -// - "rtf_mean" replaced by "rtf_wall_mean" (wall-clock RTF) and -// "rtf_compute_mean" (phase-timer sum RTF). Wall is the -// user-visible number; compute is for backward compat. -// -// The Python driver accepts both v1 and v2 schemas. +// Schema version: "transcribe-bench-v2". Reports "rtf_wall_mean" +// (wall-clock RTF, the user-visible number) and "rtf_compute_mean" +// (phase-timer sum RTF). The Python driver accepts both v1 and v2 schemas. #include "transcribe.h" #include "wav.h" @@ -198,7 +193,6 @@ int main(int argc, char ** argv) { if (args.quiet) transcribe_log_set(nullptr, nullptr); const bool quiet = args.quiet; - // Load WAV. if (!quiet) std::fprintf(stderr, "loading sample %s\n", args.sample_path.c_str()); std::vector pcm; std::string wav_err; @@ -225,7 +219,6 @@ int main(int argc, char ** argv) { const char * backend_c = transcribe_model_backend(model); const std::string backend = (backend_c && *backend_c) ? backend_c : "unknown"; - // Init context. struct transcribe_session_params cp; transcribe_session_params_init(&cp); cp.n_threads = args.n_threads; struct transcribe_session * ctx = nullptr; diff --git a/tools/transcribe-quantize/main.cpp b/tools/transcribe-quantize/main.cpp index 169f9476..1521c283 100644 --- a/tools/transcribe-quantize/main.cpp +++ b/tools/transcribe-quantize/main.cpp @@ -7,16 +7,9 @@ // llama-quantize uses), so the output is bit-compatible with anything // that consumes standard GGUF blocks. // -// Why a separate C++ tool instead of doing this in convert-.py: -// -// gguf-py 0.18 ships pure-Python implementations of F32, F16, BF16, -// Q8_0, Q4_0, Q4_1, Q5_0, Q5_1 — but every K-quant (Q4_K, Q5_K, Q6_K, -// Q8_K) and every IQ-quant raises NotImplementedError. Bridging -// ggml's C quantizers from Python via ctypes works but is fragile. -// Following whisper.cpp / llama.cpp's pattern, the conversion path -// is Python (one-shot checkpoint → source/reference-dtype GGUF), and -// every derived quant is produced by a small C++ binary that links -// libggml directly. +// Quants live in C++ rather than the Python converters because gguf-py +// has no K-quant / IQ-quant support: conversion stays Python (checkpoint → +// reference-dtype GGUF), every derived quant is produced here against libggml. // // The tool is intentionally family-aware in exactly one place: // classify_tensor(). If you change a family tensor catalog, keep this @@ -170,7 +163,7 @@ int main(int argc, char ** argv) { std::printf("transcribe-quantize: %s -> %s (preset %s)\n", in_path, out_path, preset->name); - // ---- Load input gguf with tensor data into a ggml_context ---- + // Load input gguf with tensor data into a ggml_context. // // no_alloc=false + ctx=&ctx_in lets ggml allocate one big buffer // and load every tensor's bytes into it. After this returns, every @@ -189,7 +182,7 @@ int main(int argc, char ** argv) { const int64_t n_tensors = gguf_get_n_tensors(gguf_in); std::printf("input: %lld tensors\n", (long long)n_tensors); - // ---- Plan: per-tensor target type + total output buffer size ---- + // Plan: per-tensor target type + total output buffer size. // // We need a fresh ggml_context for the output tensors that owns // their (possibly newly-quantized) data. The size is the sum of @@ -268,7 +261,7 @@ int main(int argc, char ** argv) { print_stats(out_stats, "output"); std::printf("\n"); - // ---- Allocate the output ggml_context + gguf_context ---- + // Allocate the output ggml_context + gguf_context. ggml_init_params out_init{}; out_init.mem_size = mem_size; out_init.mem_buffer = nullptr; // ggml will malloc internally @@ -298,7 +291,7 @@ int main(int argc, char ** argv) { // gguf-dump display it. gguf_set_val_u32(gguf_out, "general.file_type", preset->file_type); - // ---- Per-tensor: dequant → fp32 → requant → add to ctx_out ---- + // Per-tensor: dequant → fp32 → requant → add to ctx_out. std::vector fp32_scratch; int64_t requantized = 0; int64_t copied = 0; @@ -375,7 +368,7 @@ int main(int argc, char ** argv) { std::printf("processed: %lld requantized, %lld copied\n", (long long)requantized, (long long)copied); - // ---- Write the output gguf ---- + // Write the output gguf. if (!gguf_write_to_file(gguf_out, out_path, /*only_meta=*/false)) { std::fprintf(stderr, "transcribe-quantize: gguf_write_to_file failed\n"); ggml_free(ctx_out);