From 9e0c3c300e59fee870f19cb39af9825cb48284c6 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 25 Jun 2026 07:58:22 +0800 Subject: [PATCH 1/3] encoder and decoder opts --- src/arch/parakeet/decoder.cpp | 60 +++++++++++++++++++++++++++++------ src/arch/parakeet/encoder.cpp | 25 +++++++++++++-- 2 files changed, 74 insertions(+), 11 deletions(-) diff --git a/src/arch/parakeet/decoder.cpp b/src/arch/parakeet/decoder.cpp index 94bc9bdb..a7904ec8 100644 --- a/src/arch/parakeet/decoder.cpp +++ b/src/arch/parakeet/decoder.cpp @@ -546,8 +546,23 @@ inline void linear(const float * W, #endif for (int r = 0; r < out_dim; ++r) { const float * row = W + static_cast(r) * static_cast(in_dim); - float acc = 0.0f; - for (int c = 0; c < in_dim; ++c) { + // Eight independent accumulators break the single-acc reduction + // dependency chain so the compiler can auto-vectorize the dot (the + // serial version pins it to one FMA's latency — ~7x off on MSVC, the + // dominant decode cost). Pure fp32, DL-safe, portable: no intrinsics, + // no ggml-cpu symbols, no arch flags. The tree-shaped final sum changes + // float reassociation only (argmax-robust, transcript-identical). + float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0; + int c = 0; + const int in8 = in_dim & ~7; + for (; c < in8; c += 8) { + a0 += row[c + 0] * x[c + 0]; a1 += row[c + 1] * x[c + 1]; + a2 += row[c + 2] * x[c + 2]; a3 += row[c + 3] * x[c + 3]; + a4 += row[c + 4] * x[c + 4]; a5 += row[c + 5] * x[c + 5]; + a6 += row[c + 6] * x[c + 6]; a7 += row[c + 7] * x[c + 7]; + } + float acc = ((a0 + a1) + (a2 + a3)) + ((a4 + a5) + (a6 + a7)); + for (; c < in_dim; ++c) { acc += row[c] * x[c]; } y[r] = acc + (b != nullptr ? b[r] : 0.0f); @@ -774,7 +789,15 @@ void joint_step(const HostJoint & j, // 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. + if (transcribe::debug::enabled()) { float max_v = out_logits[0]; for (int i = 1; i < j.joint_n; ++i) { if (out_logits[i] > max_v) max_v = out_logits[i]; @@ -813,20 +836,39 @@ void precompute_enc_proj(const HostJoint & j, 1.0f, enc_out, d_enc, j.enc_w.data(), d_enc, 0.0f, out.data(), joint_h); -#else + // Add bias to every row. for (int t = 0; t < T; ++t) { - const float * frame = enc_out + static_cast(t) * static_cast(d_enc); float * proj = out.data() + static_cast(t) * static_cast(joint_h); - linear(j.enc_w.data(), frame, nullptr, joint_h, d_enc, proj, n_threads); + for (int k = 0; k < joint_h; ++k) { + proj[k] += j.enc_b[static_cast(k)]; + } } +#else + // No BLAS: this is a [T, d_enc] x [d_enc, joint_h]^T GEMM. The previous code + // ran it as T serial sgemv calls via linear(), and linear() only threads when + // out_dim >= 2048 — joint_h is below that, so the whole projection ran + // single-threaded (~90 ms for T=138 on this model, the dominant decode cost, + // paid on CPU even under the Vulkan backend since the decoder is host code). + // The rows are fully independent, so parallelize over T and fold the bias in. + // The inner dot auto-vectorizes (FMA/AVX-512) under the project's arch flags. + const float * enc_w = j.enc_w.data(); + const float * enc_b = j.enc_b.data(); +#ifdef _OPENMP + #pragma omp parallel for schedule(static) num_threads(n_threads > 0 ? n_threads : 1) #endif - // Add bias to every row. for (int t = 0; t < T; ++t) { + const float * frame = enc_out + static_cast(t) * static_cast(d_enc); float * proj = out.data() + static_cast(t) * static_cast(joint_h); - for (int k = 0; k < joint_h; ++k) { - proj[k] += j.enc_b[static_cast(k)]; + for (int r = 0; r < joint_h; ++r) { + const float * row = enc_w + static_cast(r) * static_cast(d_enc); + float acc = 0.0f; + for (int c = 0; c < d_enc; ++c) { + acc += row[c] * frame[c]; + } + proj[r] = acc + enc_b[static_cast(r)]; } } +#endif } // Argmax over a contiguous fp32 range. Returns the index of the diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index 30a1fc64..3bd04f56 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -84,6 +84,27 @@ bool detect_direct_dw_in_block(const char * 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. +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; + if (backend != nullptr && + (std::strstr(backend, "Metal") != nullptr || std::strstr(backend, "metal") != nullptr)) { + return false; + } + return true; +} + // ----- Views ------------------------------------------------------ conf::PreEncodeView to_view(const ParakeetPreEncode & pe) { @@ -284,7 +305,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, conf::ConvPolicy policy {}; policy.direct_pw = conf::detect_direct_pw(backend_name); policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = false; + policy.direct_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 @@ -757,7 +778,7 @@ EncoderBuild build_encoder_graph_streaming( conf::ConvPolicy policy {}; policy.direct_pw = conf::detect_direct_pw(backend_name); policy.direct_dw_in_block = detect_direct_dw_in_block(backend_name); - policy.direct_dw_in_pre_encode = false; + policy.direct_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. policy.causal_pre_encode = From 144d24c5b685a01cd899a07eb2c0e799e60d2594 Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 25 Jun 2026 21:47:11 +0800 Subject: [PATCH 2/3] parakeet speedups and associated scripts --- scripts/wer/remote/cache_paths.py | 7 +- scripts/wer/remote/modal_sweep.py | 15 +- scripts/wer/remote/output_paths.py | 4 +- .../run_reference_parakeet_streaming_nemo.py | 17 +- src/arch/parakeet/encoder.cpp | 102 ++++++-- src/arch/parakeet/encoder.h | 25 ++ src/arch/parakeet/model.cpp | 230 ++++++++++++++++-- src/arch/parakeet/parakeet.h | 26 ++ src/conformer/conformer.cpp | 175 ++++++++++--- src/conformer/conformer.h | 45 +++- 10 files changed, 563 insertions(+), 83 deletions(-) diff --git a/scripts/wer/remote/cache_paths.py b/scripts/wer/remote/cache_paths.py index 3e5f8a67..5093f4b9 100644 --- a/scripts/wer/remote/cache_paths.py +++ b/scripts/wer/remote/cache_paths.py @@ -15,6 +15,7 @@ def hyp_cache_paths( timestamps: str = "none", language: str = "", stream_chunk_ms: int = 0, + stream_att_right: int = -1, ) -> tuple[str, str]: """Deterministic Volume paths for the (model, dataset, subset, batch, sort) tuple. @@ -40,8 +41,12 @@ def hyp_cache_paths( # Streaming (--stream-chunk-ms) produces different hyps than the offline # path, so it gets its own cache slot. 0 = offline (no tag, compatible). stream_tag = "" if stream_chunk_ms <= 0 else f".stream{stream_chunk_ms}ms" + # Cache-aware latency preset (--stream-att-right). -1 = unset (model + # default R), no tag so it stays compatible with already-cached entries; + # any explicit R gets its own slot so e.g. R=13 and R=0 never collide. + r_tag = "" if stream_att_right < 0 else f".r{stream_att_right}" base = (f"/data/wer/hyps/{hyp_fp}/{slug}." - f"{dataset_id(dataset_spec)}.{subset_tag}{bs_tag}{sort_tag}{ts_tag}{lang_tag}{stream_tag}") + f"{dataset_id(dataset_spec)}.{subset_tag}{bs_tag}{sort_tag}{ts_tag}{lang_tag}{stream_tag}{r_tag}") return f"{base}.jsonl", f"{base}.summary.json" diff --git a/scripts/wer/remote/modal_sweep.py b/scripts/wer/remote/modal_sweep.py index 43eddb43..c1ab722f 100644 --- a/scripts/wer/remote/modal_sweep.py +++ b/scripts/wer/remote/modal_sweep.py @@ -461,6 +461,7 @@ def _run_wer_impl( timestamps: str = "none", language: str = "", stream_chunk_ms: int = 0, + stream_att_right: int = -1, dataset_status: dict | None = None, ) -> dict: """Run scripts/wer/run.py on `n_utts` (or full manifest) and return the @@ -480,7 +481,7 @@ def _run_wer_impl( # when a hyp for this (fingerprint, model, dataset, subset) already exists. cache_hyp, cache_sum = hyp_cache_paths( HYP_FP, model_file, dataset_spec, n_utts, batch_size, sort_by_length, - timestamps, language, stream_chunk_ms) + timestamps, language, stream_chunk_ms, stream_att_right) if os.path.exists(cache_hyp) and os.path.exists(cache_sum) \ and os.path.getsize(cache_hyp) > 0: _log_prepared_dataset("wer", dataset_status) @@ -537,6 +538,8 @@ def _run_wer_impl( cmd += ["--language", language] if stream_chunk_ms and stream_chunk_ms > 0: cmd += ["--stream-chunk-ms", str(stream_chunk_ms)] + if stream_att_right >= 0: + cmd += ["--stream-att-right", str(stream_att_right)] print(f"[wer] $ {' '.join(cmd)}") t0 = time.time() rc, stderr_tail = run_subprocess_capturing_stderr(cmd, cwd="/work", env=env) @@ -584,6 +587,7 @@ def _run_wer_impl( "utt_per_s": subset_count / wall_s if wall_s > 0 else 0.0, "batch_size": batch_size, "stream_chunk_ms": stream_chunk_ms, + "stream_att_right": stream_att_right, "mel_s": mel_ms / 1000.0, "encode_s": enc_ms / 1000.0, "decode_s": dec_ms / 1000.0, @@ -633,6 +637,7 @@ def runner( timestamps: str = "none", language: str = "", stream_chunk_ms: int = 0, + stream_att_right: int = -1, dataset_status: dict | None = None, ) -> dict: # Prefer the build_dir the local entrypoint computed and built into: @@ -645,6 +650,7 @@ def runner( batch_size=batch_size, sort_by_length=sort_by_length, timestamps=timestamps, language=language, stream_chunk_ms=stream_chunk_ms, + stream_att_right=stream_att_right, dataset_status=dataset_status, ) @@ -885,6 +891,7 @@ def sweep( timestamps: str = "none", language: str = "", stream_chunk_ms: int = 0, + stream_att_right: int = -1, ) -> None: """Fan WER across one or more models on one GPU class. @@ -968,7 +975,8 @@ def sweep( n = None if n_utts < 0 else n_utts futs = [(c, runner.spawn(c["repo"], c["file"], c["dataset"], n, c["bs"], sort_by_length, build_dir, timestamps, - language, stream_chunk_ms, dataset_status)) + language, stream_chunk_ms, stream_att_right, + dataset_status)) for c in cells] rows, failures = [], [] @@ -982,7 +990,8 @@ def sweep( p = write_hyp(repo_root, res["hyp_jsonl"], c["file"], c["dataset"], batch_size=(bs if bs != 1 else None), timestamps=timestamps, - stream_chunk_ms=stream_chunk_ms) + stream_chunk_ms=stream_chunk_ms, + stream_att_right=stream_att_right) s = res["summary"] rows.append((slug, c["dataset"], s["n_utts"], s["audio_s"], s["wall_s"], s["rtf_wall"], str(p))) diff --git a/scripts/wer/remote/output_paths.py b/scripts/wer/remote/output_paths.py index b3ee81e4..7f89062f 100644 --- a/scripts/wer/remote/output_paths.py +++ b/scripts/wer/remote/output_paths.py @@ -30,6 +30,7 @@ def write_hyp( batch_size: int | None = None, timestamps: str = "none", stream_chunk_ms: int = 0, + stream_att_right: int = -1, ) -> Path: out_dir = root / "reports" / "wer" out_dir.mkdir(parents=True, exist_ok=True) @@ -38,6 +39,7 @@ def write_hyp( bs_tag = "" if batch_size is None else f".b{batch_size}" ts_tag = "" if timestamps == "none" else f".ts-{timestamps}" stream_tag = "" if stream_chunk_ms <= 0 else f".stream{stream_chunk_ms}ms" - out_path = out_dir / f"{slug}.{ds}{bs_tag}{ts_tag}{stream_tag}.jsonl" + r_tag = "" if stream_att_right < 0 else f".r{stream_att_right}" + out_path = out_dir / f"{slug}.{ds}{bs_tag}{ts_tag}{stream_tag}{r_tag}.jsonl" out_path.write_text(hyp_jsonl) return out_path diff --git a/scripts/wer/run_reference_parakeet_streaming_nemo.py b/scripts/wer/run_reference_parakeet_streaming_nemo.py index 4768be6f..2b859f0e 100644 --- a/scripts/wer/run_reference_parakeet_streaming_nemo.py +++ b/scripts/wer/run_reference_parakeet_streaming_nemo.py @@ -55,6 +55,16 @@ def main() -> int: p.add_argument("--pad-and-drop-preencoded", action="store_true", help="Treat the first chunk like subsequent (ONNX-export " "first-chunk semantics).") + # Uniform reference contract (modal_sweep _run_one_reference always passes + # these). Streaming is inherently per-utterance, so --batch-size >1 and + # --mode are accepted and ignored; --device selects the torch device. + p.add_argument("--device", default="cuda", + help="Torch device (cuda|cpu); falls back to cpu if cuda " + "is unavailable.") + p.add_argument("--batch-size", type=int, default=1, + help="Uniform-contract arg; ignored (streaming is bs=1).") + p.add_argument("--mode", default="streaming", + help="Uniform-contract arg; this runner is always streaming.") args = p.parse_args() if not args.manifest.exists(): @@ -95,6 +105,9 @@ def main() -> int: else: raise last model.eval() + dev = "cuda" if (args.device == "cuda" and torch.cuda.is_available()) else "cpu" + model = model.to(dev) + print(f"device: {dev}") load_ms = (time.monotonic() - t0) * 1000 if model.encoder.att_context_style != "chunked_limited": @@ -167,8 +180,8 @@ def transcribe_one_streaming(audio_path: str) -> str: channel_len, previous_hypotheses, ) = model.conformer_stream_step( - processed_signal=chunk_audio.to(torch.float32), - processed_signal_length=chunk_lengths, + processed_signal=chunk_audio.to(dev, torch.float32), + processed_signal_length=chunk_lengths.to(dev), cache_last_channel=cache_lc, cache_last_time=cache_lt, cache_last_channel_len=channel_len, diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index 3bd04f56..e6897cca 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -865,16 +865,44 @@ 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; - // ----- Pos_emb sized for T_virtual ----- - const int64_t pos_len = 2 * T_virt - 1; - eb.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, - hp.enc_d_model, pos_len); - ggml_set_name(eb.pos_emb_in, "stream.pos_emb.in"); - ggml_set_input(eb.pos_emb_in); - - // Chunked mask: [T_virtual, T_virtual]. Built host-side. + // 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). + const bool kv_cache_mode = + !transcribe::debug::enabled() && + static_cast(cache_io.k_in.size()) == n_layers && + static_cast(cache_io.v_in.size()) == n_layers; + + // ----- Pos_emb (rectangular: pos_len = T_q + T_kv - 1) ----- + const int64_t pos_len = T_virt + T_q_new - 1; + // Rel-pos projection memoization: when the driver's cache matches + // this chunk's pos_len, every block consumes its precomputed + // projection and the graph needs no pos_emb input at all. + const bool pos_proj_hit = + cache_io.pos_proj_len == pos_len && + static_cast(cache_io.pos_proj.size()) == n_layers; + if (!pos_proj_hit) { + eb.pos_emb_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + hp.enc_d_model, pos_len); + ggml_set_name(eb.pos_emb_in, "stream.pos_emb.in"); + ggml_set_input(eb.pos_emb_in); + } + + // Chunked mask: [T_virtual, T_q_new] (keys x queries). Built host-side. eb.chunked_mask_in = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, - T_virt, T_virt, 1, 1); + T_virt, T_q_new, + 1, 1); ggml_set_name(eb.chunked_mask_in, "stream.attn.chunked_mask.in"); ggml_set_input(eb.chunked_mask_in); @@ -891,7 +919,13 @@ EncoderBuild build_encoder_graph_streaming( bparams.n_head = hp.enc_n_heads; bparams.conv_kernel = hp.enc_conv_kernel; bparams.kv_type = kv_type; - bparams.use_flash = true; + // 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. + 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 @@ -914,28 +948,60 @@ EncoderBuild build_encoder_graph_streaming( cache_io.time_out.assign(n_layers, nullptr); const int k_minus_1 = static_cast(cache_io.time_in[0]->ne[0]); + if (kv_cache_mode) { + cache_io.k_out.assign(n_layers, nullptr); + 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. - ggml_tensor * cache_ch_out = ggml_new_tensor_2d( - ctx, GGML_TYPE_F32, hp.enc_d_model, T_cache); + // KV mode swaps the channel cache for pre-projected K/V (the + // channel_out slot stays null; the driver rotates whichever + // outs are present). ggml_tensor * cache_tm_out = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, k_minus_1, hp.enc_d_model); - if (cache_ch_out == nullptr || cache_tm_out == nullptr) return eb; + if (cache_tm_out == nullptr) return eb; char nm[64]; - std::snprintf(nm, sizeof(nm), "stream.cache_ch_out.%d", i); - ggml_set_name(cache_ch_out, nm); std::snprintf(nm, sizeof(nm), "stream.cache_tm_out.%d", i); ggml_set_name(cache_tm_out, nm); - ggml_set_output(cache_ch_out); ggml_set_output(cache_tm_out); + ggml_tensor * cache_ch_out = nullptr; + ggml_tensor * cache_k_out = nullptr; + ggml_tensor * cache_v_out = nullptr; + if (kv_cache_mode) { + cache_k_out = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, hp.enc_d_model, T_cache); + cache_v_out = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, hp.enc_d_model, T_cache); + if (cache_k_out == nullptr || cache_v_out == nullptr) return eb; + std::snprintf(nm, sizeof(nm), "stream.cache_k_out.%d", i); + ggml_set_name(cache_k_out, nm); + std::snprintf(nm, sizeof(nm), "stream.cache_v_out.%d", i); + ggml_set_name(cache_v_out, nm); + ggml_set_output(cache_k_out); + ggml_set_output(cache_v_out); + } else { + cache_ch_out = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, hp.enc_d_model, T_cache); + if (cache_ch_out == nullptr) return eb; + std::snprintf(nm, sizeof(nm), "stream.cache_ch_out.%d", i); + ggml_set_name(cache_ch_out, nm); + ggml_set_output(cache_ch_out); + } + bparams.streaming_channel_in = cache_io.channel_in[i]; bparams.streaming_channel_out = cache_ch_out; bparams.streaming_time_in = cache_io.time_in[i]; bparams.streaming_time_out = cache_tm_out; + bparams.streaming_pos_proj_in = + pos_proj_hit ? cache_io.pos_proj[i] : nullptr; + bparams.streaming_kv_k_in = kv_cache_mode ? cache_io.k_in[i] : nullptr; + bparams.streaming_kv_v_in = kv_cache_mode ? cache_io.v_in[i] : nullptr; + bparams.streaming_kv_k_out = cache_k_out; + bparams.streaming_kv_v_out = cache_v_out; x = conf::build_conformer_block(ctx, x, eb.pos_emb_in, to_view(w.blocks[i]), @@ -944,6 +1010,10 @@ EncoderBuild build_encoder_graph_streaming( cache_io.channel_out[i] = cache_ch_out; cache_io.time_out[i] = cache_tm_out; + if (kv_cache_mode) { + cache_io.k_out[i] = cache_k_out; + cache_io.v_out[i] = cache_v_out; + } } // Tag the streaming per-chunk encoder output so the host-side dump diff --git a/src/arch/parakeet/encoder.h b/src/arch/parakeet/encoder.h index 4d5c68ae..afdb05f1 100644 --- a/src/arch/parakeet/encoder.h +++ b/src/arch/parakeet/encoder.h @@ -255,6 +255,31 @@ struct StreamingEncoderCacheIO { // build_encoder_graph_streaming so the caller can copy back. 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. + 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). + std::vector pos_proj; + int pos_proj_len = -1; }; // Build a cache-aware streaming encoder graph. Same overall topology diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index b2372680..7994d579 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -122,8 +122,21 @@ ParakeetSession::~ParakeetSession() { ggml_free(stream_caches.ctx); stream_caches.ctx = nullptr; } + // Rel-pos projection memo (own ctx + buffer; geometry-keyed). + if (stream_caches.pos_proj_buf != nullptr) { + ggml_backend_buffer_free(stream_caches.pos_proj_buf); + stream_caches.pos_proj_buf = nullptr; + } + if (stream_caches.pos_proj_ctx != nullptr) { + ggml_free(stream_caches.pos_proj_ctx); + stream_caches.pos_proj_ctx = nullptr; + } + stream_caches.pos_proj.clear(); + stream_caches.pos_proj_len = -1; stream_caches.last_channel.clear(); stream_caches.last_time.clear(); + stream_caches.last_k.clear(); + stream_caches.last_v.clear(); stream_caches.initialized = false; } @@ -385,10 +398,11 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, return TRANSCRIBE_ERR_INVALID_ARG; } - // 2 tensors per layer (last_channel + last_time) plus headroom - // for descriptor / view tensors created during graph build. + // 4 tensors per layer (last_channel + last_time + last_k + last_v) + // plus headroom for descriptor / view tensors created during graph + // build. const size_t ctx_size = - static_cast(2 * n_layer + 8) * ggml_tensor_overhead(); + static_cast(4 * n_layer + 8) * ggml_tensor_overhead(); ggml_init_params ip {}; ip.mem_size = ctx_size; ip.mem_buffer = nullptr; @@ -402,6 +416,8 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, pc->stream_caches.last_channel.assign(n_layer, nullptr); pc->stream_caches.last_time.assign(n_layer, nullptr); + pc->stream_caches.last_k.assign(n_layer, nullptr); + pc->stream_caches.last_v.assign(n_layer, nullptr); for (int i = 0; i < n_layer; ++i) { ggml_tensor * lc = ggml_new_tensor_2d(pc->stream_caches.ctx, GGML_TYPE_F32, @@ -409,7 +425,15 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, ggml_tensor * lt = ggml_new_tensor_2d(pc->stream_caches.ctx, GGML_TYPE_F32, k_minus_1, d_model); - if (lc == nullptr || lt == nullptr) { + ggml_tensor * lk = ggml_new_tensor_2d(pc->stream_caches.ctx, + GGML_TYPE_F32, + d_model, T_cache); + ggml_tensor * lv = ggml_new_tensor_2d(pc->stream_caches.ctx, + GGML_TYPE_F32, + d_model, T_cache); + if (lc == nullptr || lt == nullptr || + lk == nullptr || lv == nullptr) + { ggml_free(pc->stream_caches.ctx); pc->stream_caches.ctx = nullptr; return TRANSCRIBE_ERR_OOM; @@ -419,8 +443,14 @@ transcribe_status init_streaming_caches(ParakeetSession * pc, ggml_set_name(lc, name); std::snprintf(name, sizeof(name), "stream.cache.last_time.%d", i); ggml_set_name(lt, name); + std::snprintf(name, sizeof(name), "stream.cache.last_k.%d", i); + ggml_set_name(lk, name); + std::snprintf(name, sizeof(name), "stream.cache.last_v.%d", i); + ggml_set_name(lv, name); pc->stream_caches.last_channel[i] = lc; pc->stream_caches.last_time[i] = lt; + pc->stream_caches.last_k[i] = lk; + pc->stream_caches.last_v[i] = lv; } pc->stream_caches.buffer = @@ -2107,10 +2137,10 @@ static int64_t us_to_ms(int64_t us) { // pc->pos_div_term are resized in place. void fill_streaming_pos_emb(ParakeetSession * pc, ggml_tensor * pos_emb_in, - int d_model) + int d_model, + int zero_index) { const int pos_len = static_cast(pos_emb_in->ne[1]); - const int zero_index = (pos_len - 1) / 2; pc->pos_buf.assign(static_cast(pos_len) * d_model, 0.0f); pc->pos_div_term.resize(static_cast(d_model / 2)); @@ -2149,6 +2179,8 @@ void fill_streaming_chunked_mask(ggml_tensor * mask_in, int att_context_left, int att_context_right) { + const int n_q_rows = static_cast(mask_in->ne[1]); + const int q_first = T_virtual - n_q_rows; const int chunk_size = att_context_right + 1; const int left_chunks = (chunk_size > 0) ? (att_context_left / chunk_size) @@ -2156,15 +2188,15 @@ void fill_streaming_chunked_mask(ggml_tensor * mask_in, const int invalid_cache_threshold = T_cache - channel_len; std::vector mask_buf( - static_cast(T_virtual) * static_cast(T_virtual)); - for (int q = 0; q < T_virtual; ++q) { + static_cast(n_q_rows) * static_cast(T_virtual)); + for (int q = q_first; q < T_virtual; ++q) { const int q_chunk = (chunk_size > 0) ? (q / chunk_size) : 0; const int k_min_chunk = (q_chunk - left_chunks > 0) ? (q_chunk - left_chunks) : 0; const int k_min = k_min_chunk * chunk_size; const int k_max = (q_chunk + 1) * chunk_size; // exclusive float * row = mask_buf.data() + - static_cast(q) * T_virtual; + static_cast(q - q_first) * T_virtual; for (int k = 0; k < T_virtual; ++k) { const bool in_band = (k >= k_min && k < k_max); const bool cache_unfilled = (k < invalid_cache_threshold); @@ -2177,6 +2209,111 @@ void fill_streaming_chunked_mask(ggml_tensor * mask_in, 0, mask_buf.size() * sizeof(float)); } +// 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. +transcribe_status ensure_pos_proj_cache(ParakeetSession * pc, + ParakeetModel * pm, + int pos_len) +{ + const auto & hp = pm->hparams; + const int n_layers = static_cast(pm->weights.blocks.size()); + const int d_model = hp.enc_d_model; + const int n_head = hp.enc_n_heads; + const int head_dim = d_model / n_head; + + if (pc->pos_buf.size() < + static_cast(pos_len) * static_cast(d_model)) + { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet stream: pos_proj cache fill without pos_buf " + "(%zu < %d*%d)", pc->pos_buf.size(), pos_len, d_model); + return TRANSCRIBE_ERR_BACKEND; + } + + auto & sc = pc->stream_caches; + if (sc.pos_proj_buf != nullptr) { + ggml_backend_buffer_free(sc.pos_proj_buf); + sc.pos_proj_buf = nullptr; + } + if (sc.pos_proj_ctx != nullptr) { + ggml_free(sc.pos_proj_ctx); + sc.pos_proj_ctx = nullptr; + } + sc.pos_proj.assign(static_cast(n_layers), nullptr); + sc.pos_proj_len = -1; + + ggml_init_params pip {}; + pip.mem_size = (static_cast(n_layers) + 2) * + ggml_tensor_overhead(); + pip.mem_buffer = nullptr; + pip.no_alloc = true; + sc.pos_proj_ctx = ggml_init(pip); + if (sc.pos_proj_ctx == nullptr) return TRANSCRIBE_ERR_OOM; + + for (int i = 0; i < n_layers; ++i) { + ggml_tensor * t = ggml_new_tensor_3d( + sc.pos_proj_ctx, GGML_TYPE_F32, head_dim, pos_len, n_head); + if (t == nullptr) return TRANSCRIBE_ERR_OOM; + char nm[64]; + std::snprintf(nm, sizeof(nm), "stream.pos_proj.%d", i); + ggml_set_name(t, nm); + sc.pos_proj[static_cast(i)] = t; + } + sc.pos_proj_buf = ggml_backend_alloc_ctx_tensors( + sc.pos_proj_ctx, pm->plan.primary); + if (sc.pos_proj_buf == nullptr) return TRANSCRIBE_ERR_OOM; + + // One-off projection graph. + ggml_init_params gip {}; + gip.mem_size = 4 * 1024 * 1024; + gip.mem_buffer = nullptr; + gip.no_alloc = true; + ggml_context * gctx = ggml_init(gip); + if (gctx == nullptr) return TRANSCRIBE_ERR_OOM; + + ggml_cgraph * graph = ggml_new_graph_custom(gctx, 512, false); + ggml_tensor * pos_in = ggml_new_tensor_2d( + gctx, GGML_TYPE_F32, d_model, pos_len); + ggml_set_name(pos_in, "pos_proj.pos_emb.in"); + ggml_set_input(pos_in); + for (int i = 0; i < n_layers; ++i) { + ggml_tensor * p = ggml_mul_mat( + gctx, pm->weights.blocks[static_cast(i)].attn_pos_w, + pos_in); + p = ggml_reshape_4d(gctx, p, head_dim, n_head, pos_len, 1); + p = ggml_cont(gctx, ggml_permute(gctx, p, 0, 2, 1, 3)); + ggml_tensor * cpy = ggml_cpy( + gctx, p, sc.pos_proj[static_cast(i)]); + ggml_build_forward_expand(graph, cpy); + } + + transcribe_status st = TRANSCRIBE_OK; + ggml_backend_sched_reset(pc->sched); + if (!ggml_backend_sched_alloc_graph(pc->sched, graph)) { + st = TRANSCRIBE_ERR_BACKEND; + } else { + ggml_backend_tensor_set( + pos_in, pc->pos_buf.data(), 0, + static_cast(pos_len) * d_model * sizeof(float)); + if (ggml_backend_sched_graph_compute(pc->sched, graph) != + GGML_STATUS_SUCCESS) + { + st = TRANSCRIBE_ERR_BACKEND; + } + } + ggml_free(gctx); + if (st == TRANSCRIBE_OK) { + sc.pos_proj_len = pos_len; + } + return st; +} + } // namespace // Already inside namespace transcribe::parakeet { ... }. The function is @@ -2318,8 +2455,12 @@ transcribe_status emit_streaming_chunk( } StreamingEncoderCacheIO cache_io; - cache_io.channel_in = pc->stream_caches.last_channel; - cache_io.time_in = pc->stream_caches.last_time; + cache_io.channel_in = pc->stream_caches.last_channel; + cache_io.time_in = pc->stream_caches.last_time; + cache_io.k_in = pc->stream_caches.last_k; + cache_io.v_in = pc->stream_caches.last_v; + cache_io.pos_proj = pc->stream_caches.pos_proj; + cache_io.pos_proj_len = pc->stream_caches.pos_proj_len; EncoderBuild eb = build_encoder_graph_streaming( pc->compute_ctx, pm->weights, hp, @@ -2352,8 +2493,19 @@ transcribe_status emit_streaming_chunk( static_cast(n_mel_chunk_frames) * static_cast(hp.fe_num_mels) * sizeof(float)); - // Build & upload pos_emb (sized for T_virtual). - fill_streaming_pos_emb(pc, eb.pos_emb_in, hp.enc_d_model); + const int T_virtual = static_cast(eb.chunked_mask_in->ne[0]); + 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. + 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) @@ -2363,9 +2515,6 @@ transcribe_status emit_streaming_chunk( // 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. - const int T_virtual = static_cast(eb.chunked_mask_in->ne[0]); - const int T_cache = hp.enc_att_context_left; - const int T_q_new = T_virtual - T_cache; fill_streaming_chunked_mask( eb.chunked_mask_in, T_virtual, T_cache, pc->stream_caches.channel_len, @@ -2485,24 +2634,59 @@ transcribe_status emit_streaming_chunk( // (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. + const bool kv_mode = !cache_io.k_out.empty(); for (int i = 0; i < n_layers; ++i) { - if (cache_io.channel_out[i] == nullptr || cache_io.channel_out[i]->buffer == nullptr || - cache_io.time_out[i] == nullptr || cache_io.time_out[i]->buffer == nullptr) - { + ggml_tensor * ch = cache_io.channel_out[i]; + ggml_tensor * tm = cache_io.time_out[i]; + ggml_tensor * ko = kv_mode ? cache_io.k_out[i] : nullptr; + ggml_tensor * vo = kv_mode ? cache_io.v_out[i] : nullptr; + const bool attn_ok = kv_mode + ? (ko != nullptr && ko->buffer != nullptr && + vo != nullptr && vo->buffer != nullptr) + : (ch != nullptr && ch->buffer != nullptr); + if (!attn_ok || tm == nullptr || tm->buffer == nullptr) { log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, "parakeet stream: cache_out unallocated at layer %d " "(att_context_right=%d) — builder bug", i, pc->stream_caches.att_context_right); return TRANSCRIBE_ERR_BACKEND; } - ggml_backend_tensor_copy(cache_io.channel_out[i], - pc->stream_caches.last_channel[i]); - ggml_backend_tensor_copy(cache_io.time_out[i], - pc->stream_caches.last_time[i]); + if (kv_mode) { + ggml_backend_tensor_copy(ko, pc->stream_caches.last_k[i]); + ggml_backend_tensor_copy(vo, pc->stream_caches.last_v[i]); + } else { + ggml_backend_tensor_copy(ch, pc->stream_caches.last_channel[i]); + } + ggml_backend_tensor_copy(tm, pc->stream_caches.last_time[i]); } 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. + 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) { + if (const transcribe_status st = + ensure_pos_proj_cache(pc, pm, pos_len_cur); + st != TRANSCRIBE_OK) + { + log_msg(TRANSCRIBE_LOG_LEVEL_ERROR, + "parakeet stream: pos_proj cache fill failed (%s) — " + "continuing without memoization", + transcribe_status_string(st)); + } + } + } + if (dump_on) { char namebuf[128]; std::snprintf(namebuf, sizeof(namebuf), diff --git a/src/arch/parakeet/parakeet.h b/src/arch/parakeet/parakeet.h index ebcc497e..64f987b4 100644 --- a/src/arch/parakeet/parakeet.h +++ b/src/arch/parakeet/parakeet.h @@ -213,6 +213,32 @@ 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 + // 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. + ggml_context * pos_proj_ctx = nullptr; + ggml_backend_buffer_t pos_proj_buf = nullptr; + std::vector pos_proj; + int pos_proj_len = -1; + int32_t channel_len = 0; // Absolute mel-frame cursor across the whole stream — used to diff --git a/src/conformer/conformer.cpp b/src/conformer/conformer.cpp index 40d0cdda..94ae11c4 100644 --- a/src/conformer/conformer.cpp +++ b/src/conformer/conformer.cpp @@ -634,7 +634,10 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, ggml_tensor * x, ggml_tensor * pos_emb, const BlockView & b, - const BlockParams & params) + const BlockParams & params, + ggml_tensor * x_q, + ggml_tensor * k_full, + ggml_tensor * v_full) { const int d_model = params.d_model; const int n_head = params.n_head; @@ -644,8 +647,29 @@ 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)); - const int64_t T_q = x->ne[1]; - const int64_t pos_len = pos_emb->ne[1]; + // 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. + 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; + const int64_t T_kv = kv_cached ? k_full->ne[1] : x->ne[1]; + const int64_t T_q = q_sliced ? x_q->ne[1] : x->ne[1]; + // With a precomputed pos projection, pos_emb may be null (the + // caller's graph carries no pos_emb input at all). + ggml_tensor * pos_proj = params.streaming_pos_proj_in; + const int64_t pos_len = pos_proj != nullptr ? pos_proj->ne[1] + : pos_emb->ne[1]; + if (rect && pos_len != T_q + T_kv - 1) { + std::fprintf(stderr, + "conformer rel_pos_mhsa: rectangular pos_emb length " + "%lld != T_q + T_kv - 1 = %lld\n", + (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 @@ -661,7 +685,11 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, // 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). - const bool flash = use_flash; + // 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 // in the Regular style, pos_emb arrives at the smaller @@ -678,13 +706,20 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, // ----- 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. - ggml_tensor * q = ggml_mul_mat(ctx, b.attn_q_w, x); + 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); - ggml_tensor * k = ggml_mul_mat(ctx, b.attn_k_w, x); - if (b.attn_k_b != nullptr) k = ggml_add(ctx, k, b.attn_k_b); - ggml_tensor * v = ggml_mul_mat(ctx, b.attn_v_w, x); - if (b.attn_v_b != nullptr) v = ggml_add(ctx, v, b.attn_v_b); - ggml_tensor * p = ggml_mul_mat(ctx, b.attn_pos_w, pos_emb); + // KV-cache mode: keys/values arrive pre-projected (bias included). + ggml_tensor * k = k_full; + ggml_tensor * v = v_full; + if (!kv_cached) { + k = ggml_mul_mat(ctx, b.attn_k_w, x); + if (b.attn_k_b != nullptr) k = ggml_add(ctx, k, b.attn_k_b); + v = ggml_mul_mat(ctx, b.attn_v_w, x); + if (b.attn_v_b != nullptr) v = ggml_add(ctx, v, b.attn_v_b); + } + ggml_tensor * p = pos_proj == nullptr + ? 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 @@ -702,16 +737,21 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, q_u = ggml_permute(ctx, q_u, 0, 2, 1, 3); q_v = ggml_cont(ctx, ggml_permute(ctx, q_v, 0, 2, 1, 3)); - k = ggml_reshape_4d(ctx, k, head_dim, n_head, T_q, B); + k = ggml_reshape_4d(ctx, k, head_dim, n_head, T_kv, B); k = ggml_permute(ctx, k, 0, 2, 1, 3); - v = ggml_reshape_4d(ctx, v, head_dim, n_head, T_q, B); + v = ggml_reshape_4d(ctx, v, head_dim, n_head, T_kv, B); v = ggml_permute(ctx, v, 0, 2, 1, 3); // Position scores are batch-independent (pos_emb has no batch axis), so // p keeps ne[3] == 1 and broadcasts across the batch in the mul_mat. - p = ggml_reshape_4d(ctx, p, head_dim, n_head, pos_len, 1); - p = ggml_cont(ctx, ggml_permute(ctx, p, 0, 2, 1, 3)); + // The precomputed pos_proj is exactly this tensor, memoized. + if (p != nullptr) { + p = ggml_reshape_4d(ctx, p, head_dim, n_head, pos_len, 1); + p = ggml_cont(ctx, ggml_permute(ctx, p, 0, 2, 1, 3)); + } else { + p = pos_proj; + } // ----- Position mask / bias ----------------------------------- // matrix_bd = rel_shift(q_v @ p^T), truncated to [T_q, T_q]. @@ -760,9 +800,17 @@ ggml_tensor * rel_pos_mhsa(ggml_context * ctx, } } + // rel_shift generalizes to the rectangular case: with input + // [T_q + T_kv - 1, T_q] it yields out[k, q] = in[k - q + T_q - 1, q], + // i.e. row k holds the score of key k against query q for relative + // offset (T_kv - 1) - (k - q + T_q - 1) = (T_kv - T_q) + q - k — + // exactly the query-at-absolute-position (T_kv - T_q + q) semantics + // the streaming x_q path needs. The square offline case is the + // T_q == T_kv specialization. The zero column injected by the trick + // only lands at k >= T_kv, which the view below slices off. matrix_bd = rel_shift(ctx, matrix_bd); matrix_bd = ggml_view_4d(ctx, matrix_bd, - T_q, T_q, n_head, B, + T_kv, T_q, n_head, B, matrix_bd->nb[1], matrix_bd->nb[2], matrix_bd->nb[3], /*offset=*/0); // The view is non-contiguous (nb[1] stays at parent's @@ -897,29 +945,59 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, // 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). - // 2. Virtualize x_norm for the attention call: + // 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 runs on T_virtual = T_cache + T_q_new - // positions. The pos_emb and chunked_mask the caller built - // are sized for T_virtual. - // 3. Slice the attention output to the last T_q_new rows so it - // lines up with the running residual (which still has - // T_q_new frames). + // 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 + // 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). { const bool streaming = (params.streaming_channel_in != nullptr); + // KV-cache streaming: keys/values are cached pre-projected, so + // the channel cache (and its concat/write) is skipped entirely + // and x_norm stays at the new frames throughout. + const bool kv_mode = streaming && + params.streaming_kv_k_in != nullptr && + params.streaming_kv_v_in != nullptr; const int64_t T_q_new = streaming ? x->ne[1] : 0; const int64_t T_cache = streaming ? params.streaming_channel_in->ne[1] : 0; + ggml_tensor * x_q_new = nullptr; ggml_tensor * x_norm = layer_norm(ctx, x, b.norm_attn_w, b.norm_attn_b); - if (streaming) { + // Streaming cache rotation: new_cache = concat(prev[T_q_new:], + // fresh) for the common T_q_new < cache-size case, else the + // last cache-size rows of fresh. Shared by the channel cache + // and the KV cache (identical rolling-window semantics). + const auto emit_rotated_cache = [&](ggml_tensor * prev, + ggml_tensor * fresh, + ggml_tensor * out) { + const int64_t Tc = prev->ne[1]; + ggml_tensor * new_cache = nullptr; + if (T_q_new < Tc) { + ggml_tensor * prev_tail = ggml_view_2d( + ctx, prev, prev->ne[0], Tc - T_q_new, prev->nb[1], + /*offset=*/T_q_new * prev->nb[1]); + new_cache = ggml_concat(ctx, prev_tail, fresh, /*dim=*/1); + } else { + new_cache = ggml_view_2d( + ctx, fresh, fresh->ne[0], Tc, fresh->nb[1], + (T_q_new - Tc) * fresh->nb[1]); + } + ggml_tensor * cpy = ggml_cpy(ctx, new_cache, out); + ggml_build_forward_expand(params.streaming_graph, cpy); + }; + + 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 @@ -958,23 +1036,48 @@ ggml_tensor * build_conformer_block(ggml_context * ctx, ggml_build_forward_expand(params.streaming_graph, cpy); } - // Virtual-T attention: prepend the previous cache to - // x_norm and let rel_pos_mhsa run as usual. + // Virtual-T attention with query slicing: the pre-concat + // x_norm doubles as the query-only input, so attention output + // is computed for just the T_q_new new frames while K/V span + // the full virtual window (after the concat below). + x_q_new = x_norm; x_norm = ggml_concat(ctx, params.streaming_channel_in, x_norm, /*dim=*/1); } - ggml_tensor * attn_out = rel_pos_mhsa(ctx, x_norm, pos_emb, b, params); + ggml_tensor * attn_out = nullptr; + if (kv_mode) { + // Project K/V for the new frames only, prepend the cached + // window, and rotate the cache forward. + ggml_tensor * k_new = ggml_mul_mat(ctx, b.attn_k_w, x_norm); + if (b.attn_k_b != nullptr) k_new = ggml_add(ctx, k_new, b.attn_k_b); + ggml_tensor * v_new = ggml_mul_mat(ctx, b.attn_v_w, x_norm); + if (b.attn_v_b != nullptr) v_new = ggml_add(ctx, v_new, b.attn_v_b); + + ggml_tensor * k_all = ggml_concat( + ctx, params.streaming_kv_k_in, k_new, /*dim=*/1); + ggml_tensor * v_all = ggml_concat( + ctx, params.streaming_kv_v_in, v_new, /*dim=*/1); + + if (params.streaming_kv_k_out != nullptr && + params.streaming_kv_v_out != nullptr && + params.streaming_graph != nullptr) + { + emit_rotated_cache(params.streaming_kv_k_in, k_new, + params.streaming_kv_k_out); + emit_rotated_cache(params.streaming_kv_v_in, v_new, + params.streaming_kv_v_out); + } - if (streaming) { - // attn_out: [d_model, T_virtual]. Slice to last T_q_new. - const int64_t T_virt = attn_out->ne[1]; - attn_out = ggml_view_2d( - ctx, attn_out, - attn_out->ne[0], T_q_new, - attn_out->nb[1], - (T_virt - T_q_new) * attn_out->nb[1]); - attn_out = ggml_cont(ctx, attn_out); + attn_out = rel_pos_mhsa(ctx, x_norm, pos_emb, b, params, + /*x_q=*/nullptr, k_all, v_all); + } else { + // x_q_new is the new-frame query slice when streaming, or + // nullptr offline (full self-attention). rel_pos_mhsa already + // returns only the T_q_new query rows in the streaming case, + // so no output slicing is needed. + attn_out = rel_pos_mhsa(ctx, x_norm, pos_emb, b, params, + x_q_new); } x = ggml_add(ctx, x, attn_out); diff --git a/src/conformer/conformer.h b/src/conformer/conformer.h index 1c8c2d4d..db4d9bfa 100644 --- a/src/conformer/conformer.h +++ b/src/conformer/conformer.h @@ -337,6 +337,30 @@ struct BlockParams { // Ignored in offline mode (streaming_channel_in == nullptr). 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. + 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. + 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 @@ -461,11 +485,30 @@ ggml_tensor * conv_module(ggml_context * ctx, // matrix_bd before flash_attn / soft_max_ext. // // Otherwise (Regular + att_context_* == -1) -> unrestricted attention. +// x_q (optional): query-only activation [d_model, T_q, B] with T_q != +// x->ne[1]. When non-null, K/V (and the keys of the score matrix) come +// from `x` while queries come from `x_q` — the cache-aware streaming +// case where only the new frames need attention output but the cached +// frames still serve as keys/values. Requirements in this mode: +// - pos_emb length must be T_q + T_k - 1 with the zero-offset row at +// index T_k - 1 (row i holds relative offset (T_k - 1) - i) +// - attn_chunked_mask (if any) must be [T_k, T_q] +// - the flash path is bypassed (manual mul_mat + soft_max only) +// Null (default) keeps the historical self-attention behavior. +// +// k_full / v_full (optional, set together): pre-projected keys/values +// [d_model, T_k, 1, 1] (bias already applied). When set, the helper +// skips its own K/V projections entirely, queries come from `x`, and +// the same rectangular requirements as x_q apply. Streaming KV-cache +// path; mutually exclusive with x_q. ggml_tensor * rel_pos_mhsa(ggml_context * ctx, ggml_tensor * x, ggml_tensor * pos_emb, const BlockView & b, - const BlockParams & params); + const BlockParams & params, + ggml_tensor * x_q = nullptr, + ggml_tensor * k_full = nullptr, + ggml_tensor * v_full = nullptr); // =========================================================================== // Top-level block + pre-encode From 74f86e433fe36ca623a31cb3c26dbfc7c20758ac Mon Sep 17 00:00:00 2001 From: CJ Pais Date: Thu, 25 Jun 2026 21:52:32 +0800 Subject: [PATCH 3/3] doc for streaming --- .../nemotron-speech-streaming-en-0.6b.md | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/models/nemotron-speech-streaming-en-0.6b.md b/docs/models/nemotron-speech-streaming-en-0.6b.md index 38ef6da6..5083de66 100644 --- a/docs/models/nemotron-speech-streaming-en-0.6b.md +++ b/docs/models/nemotron-speech-streaming-en-0.6b.md @@ -50,6 +50,32 @@ NVIDIA's self-reported number on the same split at `att_context_size=[70, 13]` (1.12s chunk, w/o PnC) is 2.32% (from the [HF model card](https://huggingface.co/nvidia/nemotron-speech-streaming-en-0.6b)). +## Streaming WER + +In cache-aware streaming mode the encoder runs incrementally over fixed +chunks while carrying constant-memory caches. WER on the full LibriSpeech +test-clean split (2620 utterances) at the default `att_context_right=13` +setting. + +| Quantization | WER (streaming, R=13) | +| --- | ---: | +| F16 | 2.29% | +| Q8_0 | 2.31% | + +Streaming matches offline (2.31%) and NVIDIA's NeMo cache-aware streaming +reference on the same split (2.31%) + +Reproduction: + +```bash +uv run scripts/wer/run.py \ + --cli build/bin/transcribe-cli \ + --model models/nemotron-speech-streaming-en-0.6b/nemotron-speech-streaming-en-0.6b-Q8_0.gguf \ + --manifest --out hyps.jsonl \ + --stream-chunk-ms 1040 --stream-att-right 13 +uv run scripts/wer/score.py hyps.jsonl +``` + ## Quick Start ```bash